diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 97226159daeaf..10dce6b25a958 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 @@ -3213,7 +3207,7 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, StructArray, UInt32Array, UInt64Array, + Int64Array, NullArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; @@ -4398,6 +4392,15 @@ mod tests { } fn partial_reduce_test_aggregate() -> Result { + partial_reduce_test_aggregate_with_batches(1) + } + + /// Partial-reduce aggregate over `num_input_batches` identical input batches + /// of partial states, each reducing to groups `1, 2, 3` with sums + /// `50, 20, 30`. + fn partial_reduce_test_aggregate_with_batches( + num_input_batches: usize, + ) -> Result { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Float64, false), @@ -4430,7 +4433,7 @@ mod tests { ], )?; let partial_reduce_input = TestMemoryExec::try_new_exec( - &[vec![partial_state_batch]], + &[vec![partial_state_batch; num_input_batches]], Arc::clone(&partial_schema), None, )?; @@ -4466,26 +4469,162 @@ mod tests { Ok(()) } - /// Spilling behavior is not implemented for partial-reduce stream yet, so fall - /// back to the existing `GroupedHashAggregateStream` + /// Partial-reduce hash aggregation emits its accumulated partial states early + /// under memory pressure instead of failing, and the early-emitted states + /// still merge into the correct result. #[tokio::test] - async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { - let partial_reduce = partial_reduce_test_aggregate()?; + async fn partial_reduce_aggregate_with_memory_limit_emits_early() -> Result<()> { + let num_input_batches = 3; + let partial_reduce = + partial_reduce_test_aggregate_with_batches(num_input_batches)?; let runtime = RuntimeEnvBuilder::new() .with_memory_limit(1, 1.0) .build_arc()?; - let task_ctx = - Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().set_bool( - "datafusion.execution.enable_migration_aggregate", - true, - )) - .with_runtime(runtime), - ); + // A batch size smaller than the number of flushed groups also covers + // splitting one flush across several output batches. + let batch_size = 2; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)) + .with_runtime(runtime), + ); 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 output = collect(stream).await?; + + // The table is flushed after every input batch, so each of the three + // groups is emitted once per input batch instead of being merged into a + // single row. Each flush is sliced into batches of 2 and 1 rows. + assert_eq!(output.len(), 2 * num_input_batches); + assert_snapshot!(batches_to_string(&output), @r" + +---+-------------+ + | a | SUM(b)[sum] | + +---+-------------+ + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + +---+-------------+ + "); + + Ok(()) + } + + /// Same shape as [`partial_reduce_test_aggregate_with_batches`], but with multiple + /// group keys. + fn partial_reduce_test_aggregate_rows_multi_group_keys( + num_input_batches: usize, + ) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("n", DataType::Null, true), + Field::new("b", DataType::Float64, false), + ])); + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema)?, "a".to_string()), + (col("n", &schema)?, "n".to_string()), + ]); + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), + Arc::new(NullArray::new(4)), + Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let partial_reduce_input = TestMemoryExec::try_new_exec( + &[vec![partial_state_batch; num_input_batches]], + Arc::clone(&partial_schema), + None, + )?; + + AggregateExec::try_new( + AggregateMode::PartialReduce, + group_by, + aggregates, + vec![None], + partial_reduce_input, + partial_schema, + ) + } + + #[tokio::test] + async fn partial_reduce_aggregate_with_memory_limit_emits_early_multi_group_keys() + -> Result<()> { + let num_input_batches = 3; + let partial_reduce = + partial_reduce_test_aggregate_rows_multi_group_keys(num_input_batches)?; + + // Pin the representation: this is exactly the condition + // `new_group_values` uses to pick `GroupValuesRows` over + // `GroupValuesColumn`. If a `Null` `GroupColumn` is ever added, this + // assertion fires and the test stops covering the row-encoded path. + let group_schema = partial_reduce + .group_by + .group_schema(&partial_reduce.schema())?; + assert!( + !group_values::multi_group_by::supported_schema(&group_schema), + "expected the Null group column to force the GroupValuesRows fallback" + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(1, 1.0) + .build_arc()?; + let batch_size = 2; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)) + .with_runtime(runtime), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + + // Same flush cadence as the column-backed test: one flush per input + // batch, each sliced into batches of 2 and 1 rows. + assert_eq!(output.len(), 2 * num_input_batches); + assert_snapshot!(batches_to_string(&output), @r" + +---+---+-------------+ + | a | n | SUM(b)[sum] | + +---+---+-------------+ + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + +---+---+-------------+ + "); 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..858830f8390ce 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,15 @@ 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 the memory reservation cannot grow after aggregating an input batch, all +/// accumulated partial states are emitted immediately, and the remaining input +/// is aggregated with an empty table. This repeats until the input ends. +/// +/// See [`crate::aggregates::AggregateMode::PartialReduce`] for why it's allowed +/// to emit the same group multiple times. pub(crate) struct PartialReduceHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -75,6 +84,9 @@ pub(crate) struct PartialReduceHashAggregateStream { /// Input batches containing partial aggregate state rows. input: SendableRecordBatchStream, + /// Target output batch size from configuration. + batch_size: usize, + /// Execution metrics shared with the aggregate plan node. baseline_metrics: BaselineMetrics, @@ -93,10 +105,22 @@ enum PartialReduceHashAggregateState { ReadingInput { hash_table: AggregateHashTable, }, + /// A fully materialized partial-state batch being emitted incrementally + /// because the table ran out of memory while reading input. + EmittingOnMemoryPressure { + hash_table: AggregateHashTable, + // After each incremental emitting step, `remaining_groups` is updated + // with batch slicing. + remaining_groups: RecordBatch, + }, ProducingOutput { 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>>; @@ -111,28 +135,34 @@ type PartialReduceHashAggregateStateTransition = ControlFlow< impl PartialReduceHashAggregateState { fn hash_table(&self) -> &AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } fn hash_table_mut(&mut self) -> &mut AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } fn into_hash_table(self) -> AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } @@ -178,19 +208,25 @@ impl PartialReduceHashAggregateStream { Ok(Self { schema, input, + batch_size, baseline_metrics, reservation, state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), }) } - 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,46 +255,126 @@ 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 let Err(e) = self - .reservation - .try_resize(original_state.hash_table().memory_size()) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - ControlFlow::Continue(original_state) - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + // Update the memory reservation. If OOM, do early emit. + self.resize_or_emit_early(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), } } } } + /// Update the memory reservation. If the reservation succeeds, continue reading + /// input. If OOM, clear the aggregated states in the hash table, and early emit + /// them immediately. + /// + /// Returns the next state; the caller finishes the intended task based on it. + /// + /// The reservation is left at its pre-emission size while the states are being + /// emitted, because the cleared states are still held in memory as + /// `remaining_groups`. [`Self::handle_emitting_on_memory_pressure`] updates the + /// reservation once the last slice has been emitted. + /// + /// # Implementation Note + /// All accumulated states are materialized at once, and then sliced into + /// `batch_size` output batches. Emit them incrementally after blocked state + /// management is ready. + /// + /// Issue: + fn resize_or_emit_early( + &mut self, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); // Stop on drop + let resize_result = self + .reservation + .try_resize(original_state.hash_table().memory_size()); + + let oom = match resize_result { + Ok(()) => return ControlFlow::Continue(original_state), + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Self::break_with_err(e), + }; + + let state_batch_result = original_state.hash_table_mut().take_state_batch(); + + match state_batch_result { + Ok(Some(remaining_groups)) => ControlFlow::Continue( + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table: original_state.into_hash_table(), + remaining_groups, + }, + ), + // No accumulated group to emit, so early emission cannot release any + // memory: report the original error. + Ok(None) => Self::break_with_err(oom), + Err(e) => Self::break_with_err(e), + } + } + + /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state + /// batch in `batch_size`(from configuration) slices. After all slices are + /// emitted, update the memory reservation and resume reading input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_emitting_on_memory_pressure( + &mut self, + original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + let PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: batch, + } = original_state + else { + unreachable!("expected the EmittingOnMemoryPressure state") + }; + + let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + // Go back to `ReadingInput` + ( + batch, + PartialReduceHashAggregateState::ReadingInput { hash_table }, + ) + } else { + // More batches to output, continue in the current state. + let remaining = + batch.slice(self.batch_size, batch.num_rows() - self.batch_size); + let output = batch.slice(0, self.batch_size); + ( + output, + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: remaining, + }, + ) + }; + + debug_assert!(output_batch.num_rows() > 0); + ControlFlow::Break(( + Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + /// Handle ProducingOutput state - emit merged partial aggregate state batches. /// /// See comments at `poll_next()` for details. @@ -281,6 +397,8 @@ impl PartialReduceHashAggregateStream { match result { Ok(Some(batch)) => { + // The output is already materialized, so a failed resize cannot + // be acted on: keep the reservation as is and finish the output. let _ = self .reservation .try_resize(original_state.hash_table().memory_size()); @@ -300,7 +418,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), } } } @@ -325,10 +443,22 @@ impl Stream for PartialReduceHashAggregateStream { /// Aggregate one partial-state input batch, update the inner aggregate /// hash table, and continue with the next input batch. /// + /// -> EmittingOnMemoryPressure + /// The table cannot reserve enough memory. Materialize all accumulated + /// partial states and begin emitting them incrementally. + /// /// -> ProducingOutput /// Input was exhausted. Move to the next state to start outputting /// merged partial aggregate states. /// + /// EmittingOnMemoryPressure + /// -> EmittingOnMemoryPressure + /// One batch-sized slice was yielded; repeat until all materialized + /// partial states are emitted. + /// + /// -> ReadingInput + /// The materialized states were emitted; continue with the empty table. + /// /// ProducingOutput /// -> ProducingOutput /// One merged partial-state output batch was yielded; repeat to @@ -337,6 +467,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) /// ``` @@ -354,9 +491,18 @@ impl Stream for PartialReduceHashAggregateStream { state @ PartialReduceHashAggregateState::ReadingInput { .. } => { self.handle_reading_input(cx, state) } + state @ PartialReduceHashAggregateState::EmittingOnMemoryPressure { + .. + } => self.handle_emitting_on_memory_pressure(state), 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 +515,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;