From 961ee8ea6fb4efb86ea3ea2d0a7dbdc7b18cb73c Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Tue, 11 Aug 2026 17:39:22 +0800 Subject: [PATCH] refactor: return error when partial-reduce aggregation runs out of memory --- .../physical-plan/src/aggregates/mod.rs | 22 ++--- .../src/aggregates/partial_reduce_stream.rs | 96 +++++++++++++------ .../library-user-guide/upgrading/56.0.0.md | 32 +++++++ .../library-user-guide/upgrading/index.rst | 1 + 4 files changed, 113 insertions(+), 38 deletions(-) create mode 100644 docs/source/library-user-guide/upgrading/56.0.0.md diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 97226159daeaf..d720a08f1b27c 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -187,7 +187,6 @@ use datafusion_common::{ assert_eq_or_internal_err, internal_err, not_impl_err, }; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::MemoryLimit; use datafusion_expr::{Accumulator, Aggregate}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -1281,12 +1280,7 @@ impl AggregateExec { && self.group_by.is_single() } - fn should_use_partial_reduce_hash_stream(&self, context: &TaskContext) -> bool { - // TODO: implement memory-limited path and remove this limitation - if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { - return false; - } - + fn should_use_partial_reduce_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::PartialReduce && self.limit_options.is_none() && self.input_order_mode == InputOrderMode::Linear @@ -4466,10 +4460,10 @@ mod tests { Ok(()) } - /// Spilling behavior is not implemented for partial-reduce stream yet, so fall - /// back to the existing `GroupedHashAggregateStream` + /// Partial-reduce hash aggregation returns `ResourcesExhausted` when its + /// reservation cannot grow. #[tokio::test] - async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { + async fn partial_reduce_aggregate_with_memory_limit_returns_oom() -> Result<()> { let partial_reduce = partial_reduce_test_aggregate()?; let runtime = RuntimeEnvBuilder::new() .with_memory_limit(1, 1.0) @@ -4485,7 +4479,13 @@ mod tests { ); let stream = partial_reduce.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::GroupedHash(_))); + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let err = collect(stream).await.unwrap_err(); + assert!( + matches!(err.find_root(), DataFusionError::ResourcesExhausted(_)), + "expected ResourcesExhausted, got: {err}" + ); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 2f4535e66f4ef..d6dcf19368b34 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -28,7 +28,7 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; @@ -68,6 +68,12 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// This stage is useful for tree-reduce plans. It consumes the same schema as /// a final aggregate stage, but emits the same schema as a partial aggregate /// stage. +/// +/// # Memory Management +/// +/// If memory usage exceeds the budget, the stream returns an execution error. +/// +/// Larger-than-memory execution is left for future work. pub(crate) struct PartialReduceHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -97,6 +103,10 @@ enum PartialReduceHashAggregateState { hash_table: AggregateHashTable, }, Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, } type PartialReduceHashAggregatePoll = Poll>>; @@ -114,7 +124,9 @@ impl PartialReduceHashAggregateState { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table } - Self::Done => unreachable!("Done state does not hold a hash table"), + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") + } } } @@ -123,7 +135,9 @@ impl PartialReduceHashAggregateState { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table } - Self::Done => unreachable!("Done state does not hold a hash table"), + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") + } } } @@ -132,7 +146,9 @@ impl PartialReduceHashAggregateState { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table } - Self::Done => unreachable!("Done state does not hold a hash table"), + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") + } } } @@ -184,13 +200,18 @@ impl PartialReduceHashAggregateStream { }) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { + fn close_input(&mut self) { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - hash_table.start_output() + } + + fn break_with_err( + error: DataFusionError, + ) -> PartialReduceHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialReduceHashAggregateState::Error, + )) } /// Handle ReadingInput state - aggregate partial state batches into the hash table. @@ -219,41 +240,33 @@ impl PartialReduceHashAggregateStream { timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } + // If OOM, return execution error. if let Err(e) = self .reservation .try_resize(original_state.hash_table().memory_size()) { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } ControlFlow::Continue(original_state) } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), // Input ends, move to output state Poll::Ready(None) => { + self.close_input(); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); + let result = original_state.hash_table_mut().start_output(); timer.done(); match result { Ok(()) => { ControlFlow::Continue(original_state.into_producing_output()) } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Err(e) => Self::break_with_err(e), } } } @@ -281,9 +294,12 @@ impl PartialReduceHashAggregateStream { match result { Ok(Some(batch)) => { - let _ = self + if let Err(e) = self .reservation - .try_resize(original_state.hash_table().memory_size()); + .try_resize(original_state.hash_table().memory_size()) + { + return Self::break_with_err(e); + } debug_assert!(batch.num_rows() > 0); let next_state = if original_state.hash_table().is_done() { original_state.into_done() @@ -300,7 +316,7 @@ impl PartialReduceHashAggregateStream { let _ = self.reservation.try_resize(0); ControlFlow::Continue(original_state.into_done()) } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + Err(e) => Self::break_with_err(e), } } } @@ -337,6 +353,13 @@ impl Stream for PartialReduceHashAggregateStream { /// -> Done /// All merged partial-state output was emitted. /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// /// Done /// -> (end) /// ``` @@ -357,6 +380,12 @@ impl Stream for PartialReduceHashAggregateStream { state @ PartialReduceHashAggregateState::ProducingOutput { .. } => { self.handle_producing_output(state) } + state @ PartialReduceHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } state @ PartialReduceHashAggregateState::Done => { let _ = self.reservation.try_resize(0); self.state = Some(state); @@ -369,6 +398,19 @@ impl Stream for PartialReduceHashAggregateStream { self.state = Some(next_state); continue; } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!( + next_state, + PartialReduceHashAggregateState::Error + )); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(PartialReduceHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } ControlFlow::Break((poll, next_state)) => { self.state = Some(next_state); return poll; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md new file mode 100644 index 0000000000000..68a17244eb281 --- /dev/null +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -0,0 +1,32 @@ + + +# Upgrade Guides + +## DataFusion 56.0.0 + +**Note:** DataFusion `56.0.0` has not been released yet. The information provided +in this section pertains to features and changes that have already been merged +to the main branch and are awaiting release in this version. + +### `PartialReduce` aggregation no longer spills + +`AggregateMode::PartialReduce` previously used a spill-capable aggregation path +when execution memory was bounded. It now returns an execution error when it runs out +of memory space. Larger-than-memory execution for `PartialReduce` is left for future work. diff --git a/docs/source/library-user-guide/upgrading/index.rst b/docs/source/library-user-guide/upgrading/index.rst index 51c7f1413172b..870242c82a4d6 100644 --- a/docs/source/library-user-guide/upgrading/index.rst +++ b/docs/source/library-user-guide/upgrading/index.rst @@ -21,6 +21,7 @@ Upgrade Guides .. toctree:: :maxdepth: 1 + DataFusion 56.0.0 <56.0.0> DataFusion 55.0.0 <55.0.0> DataFusion 54.0.0 <54.0.0> DataFusion 53.0.0 <53.0.0>