From b32eeebf546ca14bd4f1ae7e96726bf53601dc32 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Sun, 30 Aug 2026 13:54:49 +0800 Subject: [PATCH 1/4] refactor(hash-aggr): Early emit partially aggregated result when partial-reduce mode aggregation OOM --- .../physical-plan/src/aggregates/mod.rs | 72 +++-- .../src/aggregates/partial_reduce_stream.rs | 248 +++++++++++++++--- 2 files changed, 260 insertions(+), 60 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 97226159daeaf..e070210046792 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 @@ -3217,6 +3211,7 @@ mod tests { }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; + use datafusion_common::cast::{as_float64_array, as_uint32_array}; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{DataFusionError, internal_err}; use datafusion_execution::config::SessionConfig; @@ -4398,6 +4393,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 +4434,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 +4470,50 @@ 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 emitted as batches of 2 and 1 rows. + let total_rows: usize = output.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(total_rows, 3 * num_input_batches); + assert_eq!(output.len(), 2 * num_input_batches); + + // The repeated partial states still merge into the expected sums. + let mut sums: HashMap = HashMap::new(); + for batch in &output { + let groups = as_uint32_array(batch.column(0))?; + let states = as_float64_array(batch.column(1))?; + for row in 0..batch.num_rows() { + *sums.entry(groups.value(row)).or_default() += states.value(row); + } + } + let mut sums = sums.into_iter().collect::>(); + sums.sort_by_key(|(group, _)| *group); + assert_eq!(sums, vec![(1, 150.0), (2, 60.0), (3, 90.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..9c2ae83216fb4 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, internal_datafusion_err}; 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,139 @@ impl PartialReduceHashAggregateStream { timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - 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)) + // 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 { + // Last batch to output: the emitted states are no longer held by this + // stream, so the reservation can now drop to the emptied table's size. + let table_size = hash_table.memory_size(); + let reserved = self.reservation.size(); + if let Err(e) = self.reservation.try_resize(table_size) { + // The reservation only shrinks here, which cannot fail. + return Self::break_with_err(internal_datafusion_err!( + "Partial-reduce hash aggregate failed to update its memory \ + reservation ({reserved} bytes) to the emptied table size \ + ({table_size} bytes) after early emission: {e}" + )); + } + + // 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 +410,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 +431,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 +456,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 +480,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 +504,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 +528,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; From 969bfc98fdf519bfcaa55ddd2320630d32a6944b Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Sun, 30 Aug 2026 14:08:31 +0800 Subject: [PATCH 2/4] simplify test) --- .../physical-plan/src/aggregates/mod.rs | 33 +++++++++---------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index e070210046792..97867ebaa0924 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3211,7 +3211,6 @@ mod tests { }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; - use datafusion_common::cast::{as_float64_array, as_uint32_array}; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{DataFusionError, internal_err}; use datafusion_execution::config::SessionConfig; @@ -4497,23 +4496,23 @@ mod tests { // 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 emitted as batches of 2 and 1 rows. - let total_rows: usize = output.iter().map(RecordBatch::num_rows).sum(); - assert_eq!(total_rows, 3 * num_input_batches); + // single row. Each flush is sliced into batches of 2 and 1 rows. assert_eq!(output.len(), 2 * num_input_batches); - - // The repeated partial states still merge into the expected sums. - let mut sums: HashMap = HashMap::new(); - for batch in &output { - let groups = as_uint32_array(batch.column(0))?; - let states = as_float64_array(batch.column(1))?; - for row in 0..batch.num_rows() { - *sums.entry(groups.value(row)).or_default() += states.value(row); - } - } - let mut sums = sums.into_iter().collect::>(); - sums.sort_by_key(|(group, _)| *group); - assert_eq!(sums, vec![(1, 150.0), (2, 60.0), (3, 90.0)]); + 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(()) } From 506f4e98a047992183ef9f041924482ca72ce10e Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Mon, 31 Aug 2026 22:15:02 +0800 Subject: [PATCH 3/4] review --- .../physical-plan/src/aggregates/mod.rs | 114 +++++++++++++++++- .../src/aggregates/partial_reduce_stream.rs | 15 +-- 2 files changed, 114 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 97867ebaa0924..101396b8dd4d2 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3207,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; @@ -4517,6 +4517,118 @@ mod tests { 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_group_values(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(()) + } + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. #[tokio::test] async fn ordered_partial_aggregate_planning() -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 9c2ae83216fb4..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::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; @@ -349,19 +349,6 @@ impl PartialReduceHashAggregateStream { }; let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { - // Last batch to output: the emitted states are no longer held by this - // stream, so the reservation can now drop to the emptied table's size. - let table_size = hash_table.memory_size(); - let reserved = self.reservation.size(); - if let Err(e) = self.reservation.try_resize(table_size) { - // The reservation only shrinks here, which cannot fail. - return Self::break_with_err(internal_datafusion_err!( - "Partial-reduce hash aggregate failed to update its memory \ - reservation ({reserved} bytes) to the emptied table size \ - ({table_size} bytes) after early emission: {e}" - )); - } - // Go back to `ReadingInput` ( batch, From 2bb166f8a87c57ffe953e33662110dcab9802efa Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Mon, 31 Aug 2026 22:49:55 +0800 Subject: [PATCH 4/4] fix --- datafusion/physical-plan/src/aggregates/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 101396b8dd4d2..10dce6b25a958 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4578,7 +4578,7 @@ mod tests { -> Result<()> { let num_input_batches = 3; let partial_reduce = - partial_reduce_test_aggregate_rows_group_values(num_input_batches)?; + 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