From ed05fcf56c01e16af62675682a778e3809e884e7 Mon Sep 17 00:00:00 2001 From: Yash Verma Date: Thu, 5 Nov 2026 16:51:40 +0530 Subject: [PATCH 1/4] fix: reserve incoming aggregate state before merging --- .../src/aggregates/aggregate_stream.rs | 93 ++++++++++++- .../src/aggregates/clustered_final_stream.rs | 36 ++++- .../src/aggregates/hash_stream.rs | 125 +++++++++++++++++- .../physical-plan/src/aggregates/mod.rs | 15 ++- .../cases/aggregate_memory_call_sequence.rs | 57 ++++++-- 5 files changed, 295 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 9c29f39427805..8cfa8457ce803 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -20,6 +20,7 @@ use crate::{RecordBatchStream, SendableRecordBatchStream}; use arrow::array::ArrayRef; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{Result, ScalarValue, internal_datafusion_err, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::Operator; @@ -386,6 +387,18 @@ impl AggregateStream { loop { let result = match this.input.next().await { Some(Ok(batch)) => { + // Partial states can own the entire input. Reserve the batch + // before merging it, then transfer that reservation to the + // accumulator's state below. + let previous_size = this.reservation.size(); + if this.mode.input_mode() == AggregateInputMode::Partial + && let Err(e) = this + .reservation + .try_grow(get_record_batch_memory_size(&batch)) + { + this.finished = true; + return Some((Err(e), this)); + } let result = { let elapsed_compute = this.baseline_metrics.elapsed_compute(); let _timer = elapsed_compute.timer(); // Stops on drop @@ -405,12 +418,11 @@ impl AggregateStream { Ok(allocated) }); - // allocate memory - // This happens AFTER we actually used the memory, but simplifies the whole accounting and we are OK with - // overshooting a bit. Also this means we either store the whole record batch or not. - match result - .and_then(|allocated| this.reservation.try_grow(allocated)) - { + // Once the batch is merged, its buffers are either part of + // the accumulator state or no longer retained. + match result.and_then(|allocated| { + this.reservation.try_resize(previous_size + allocated) + }) { Ok(_) => continue, Err(e) => Err(e), } @@ -596,8 +608,11 @@ mod tests { use crate::metrics::{MetricValue, MetricsSet}; use crate::test::TestMemoryExec; use crate::{ExecutionPlan, collect}; - use arrow::array::{Float64Array, Int32Array, UInt32Array}; + use arrow::array::{Float64Array, Int32Array, Int64Array, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::DataFusionError; + use datafusion_execution::memory_pool::GreedyMemoryPool; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::{ array_agg::array_agg_udaf, min_max::{max_udaf, min_udaf}, @@ -1029,4 +1044,68 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn final_aggregate_reserves_incoming_state_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let aggregate = sum_aggregate(&schema, "v", "SUM(v)")?; + let group_by = PhysicalGroupBy::default(); + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![None], + empty, + Arc::clone(&schema), + )?; + // One batch may combine many partial states. The final SUM retains + // only one number, so reserving its accumulator alone misses the batch. + let state_schema = partial.schema(); + let batch = RecordBatch::try_new( + Arc::clone(&state_schema), + vec![Arc::new(Int64Array::from(vec![1_i64; 16_384]))], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::clone(&state_schema), + None, + )?; + let final_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggregate], + vec![None], + input, + state_schema, + )?); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(64 * 1024))) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + let err = collect( + Arc::clone(&final_aggregate) as Arc, + context, + ) + .await + .unwrap_err(); + assert!(matches!( + err.find_root(), + DataFusionError::ResourcesExhausted(_) + )); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024 * 1024))) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + let output = collect(final_aggregate as Arc, context).await?; + let result = output[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.value(0), 16_384); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 6eeb37a475ffe..26650857df324 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; @@ -323,11 +324,27 @@ impl Aggregating { while let Some(batch) = self.input.next().await.transpose()? { let timer = elapsed_compute.timer(); + let batch_memory = get_record_batch_memory_size(&batch); + match reservation.try_grow(batch_memory) { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(spill_context) = self.spill_context.as_mut() else { + return Err(e); + }; + if self.table.is_empty() { + return Err(e); + } + spill_context.sort_and_spill(self.table.take_state_batch()?)?; + reservation.try_resize(self.table.memory_size())?; + reservation.try_grow(batch_memory)?; + } + Err(e) => return Err(e), + } self.table.aggregate_batch(&batch)?; - // Check the reservation before taking completed groups, as the partial - // stream does, so spilling is unchanged. Without spilling, the refusal - // is returned below only if no group is complete. + // Replace the input-batch reservation with the table footprint before + // taking completed groups, as the partial stream does. A table growth + // refusal can still be resolved by emitting completed groups. let oom = match reservation.try_resize(self.reservation_size()) { Ok(()) => None, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { @@ -768,4 +785,17 @@ mod tests { let limit = held_memory(&plan).await?; assert_resources_exhausted_under_limit(&plan, limit).await } + + #[tokio::test] + async fn clustered_final_reserves_incoming_state_batch() -> Result<()> { + // Many partial SUM states for one key merge to a small final state. + let plan = sum_over_sorted_input( + AggregateMode::Final, + &["k"], + vec![vec![vec![1; 16_384]]], + )?; + assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); + assert_resources_exhausted_under_limit(&plan, 64 * 1024).await?; + assert_output_unchanged_under_limit(&plan, 1024 * 1024).await + } } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index ca85ccb59f6bb..3e64d8e6956e5 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -12,6 +12,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{ DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, }; @@ -680,6 +681,32 @@ impl FinalHashAggregateStream { while let Some(batch) = self.input.next().await.transpose()? { let _timer = elapsed_compute.timer(); + // Partial states may retain input-sized buffers. Account for the + // incoming batch before merging it into the table. + let batch_memory = get_record_batch_memory_size(&batch); + match self.reservation.try_grow(batch_memory) { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + // Spill the current table to make room for this batch, then + // retry admission. An individual batch that still cannot fit + // must fail rather than proceed without a reservation. + let Some(spill_context) = spill_context.as_mut() else { + return Err(e.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + if hash_table.building_group_count() == 0 { + return Err(e); + } + let result = hash_table + .take_state_batch() + .and_then(|state| spill_context.sort_and_spill(state)); + self.reservation.try_resize(hash_table.memory_size())?; + result?; + self.reservation.try_grow(batch_memory)?; + } + Err(e) => return Err(e), + } hash_table.aggregate_batch(&batch)?; // Soft group limits are usually small and rarely coincide with @@ -688,11 +715,10 @@ impl FinalHashAggregateStream { let spilled = spill_context .as_ref() .is_some_and(|context| context.has_spills()); - if self.hit_soft_group_limit(hash_table) && !spilled { - break; - } + let hit_soft_limit = self.hit_soft_group_limit(hash_table) && !spilled; - // Check memory reservation, and potentially spill. + // Replace the input-batch reservation with the table footprint, + // spilling if the table cannot be reserved. let resize_result = self.reservation .try_resize(Self::reservation_size_for_table( @@ -701,7 +727,11 @@ impl FinalHashAggregateStream { )); match resize_result { - Ok(()) => {} + Ok(()) => { + if hit_soft_limit { + break; + } + } // The table cannot reserve enough memory. // Move all current states into one fully group-key-sorted spill run. @@ -1735,4 +1765,89 @@ mod tests { ); Ok(()) } + + #[tokio::test] + async fn final_hash_reserves_incoming_state_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]); + let aggregate = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + ); + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![None], + empty, + Arc::clone(&schema), + )?; + let state_schema = partial.schema(); + let batch = RecordBatch::try_new( + Arc::clone(&state_schema), + vec![ + Arc::new(Int64Array::from(vec![1_i64; 16_384])), + Arc::new(Int64Array::from(vec![1_i64; 16_384])), + ], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::clone(&state_schema), + None, + )?; + let final_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggregate], + vec![None], + input, + state_schema, + )?); + assert_eq!( + final_aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + + let config = SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(64 * 1024))) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(config.clone()) + .with_runtime(runtime), + ); + let err = collect(final_aggregate.execute(0, context)?) + .await + .unwrap_err(); + assert!(matches!( + err.find_root(), + DataFusionError::ResourcesExhausted(_) + )); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024 * 1024))) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime(runtime), + ); + let output = collect(final_aggregate.execute(0, context)?).await?; + let result = output[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.value(0), 16_384); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0943c780e117b..733e2fca656c6 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4275,9 +4275,18 @@ mod tests { assert!(final_stats.total_byte_size.get_value().is_some()); let task_ctx = if spill { - // The fair pool is shared by the partial and final tables, give each - // one its initial table plus a little, so the final aggregation spills - new_spill_ctx(2, initial_reservation(&merged_aggregate)? + 2 * 64) + // The final table must be able to reserve one incoming partial-state + // batch as well as its initial footprint. Keep only a little extra + // room so growth of the table still triggers spilling. + let max_state_batch = result + .iter() + .map(datafusion_common::utils::memory::get_record_batch_memory_size) + .max() + .unwrap_or(0); + new_spill_ctx( + 2, + initial_reservation(&merged_aggregate)? + 2 * (max_state_batch + 64), + ) } else { Arc::clone(&task_ctx) }; diff --git a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs index 853cbabf3ce50..6e2ea4dba710b 100644 --- a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs +++ b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs @@ -446,16 +446,16 @@ fn hide_byte_counts() -> Settings { /// Runs `case` as Final and as Single and returns the Final run. /// -/// The Final and Single streams account in the same way, so this asserts that -/// both make the same pool calls, with the same byte counts, once the stream -/// name is replaced, and the test snapshots only the Final run. +/// Final reserves each incoming partial-state batch before merging it, while +/// Single consumes raw rows. Their pool calls differ, but their results must +/// still match. The test snapshots the Final run. async fn run_final_and_single(case: impl Fn(AggregateMode) -> Case) -> Result { let final_run = case(AggregateMode::Final).run().await?; let single_run = case(AggregateMode::Single).run().await?; assert_eq!( - single_run.snapshot.replace("Single", "Final"), - final_run.snapshot, - "Final and Single make different pool calls" + datafusion_common::test_util::batches_to_sort_string(&single_run.output), + datafusion_common::test_util::batches_to_sort_string(&final_run.output), + "Final and Single return different results" ); Ok(final_run) } @@ -506,15 +506,21 @@ async fn spills_twice_then_replays() -> Result<()> { .await?; assert_eq!(run.clustering, GroupClusteringMode::None); assert_eq!(run.spill_count(), 2); - assert_snapshot!(run.snapshot, @r" + assert_snapshot!(run.snapshot, @" FinalHashAggregateStream[0] register can_spill=true FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] -> refused FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] -> refused FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] shrink [bytes] @@ -526,21 +532,27 @@ async fn spills_twice_then_replays() -> Result<()> { FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] shrink [bytes] @@ -600,10 +612,11 @@ async fn output_materialization_refused() -> Result<()> { let run = padded(AggregateMode::Final).run().await?; assert_eq!(run.spill_count(), 1); assert_eq!(run.output_batch_rows, vec![64; 4]); - assert_snapshot!(run.snapshot, @r" + assert_snapshot!(run.snapshot, @" FinalHashAggregateStream[0] register can_spill=true FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] + FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] -> refused FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] unregister @@ -718,17 +731,21 @@ async fn clustered_completed_groups() -> Result<()> { assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); assert_eq!(run.spill_count(), 0); assert_eq!(run.max_output_rows(), 64); - assert_snapshot!(run.snapshot, @r" + assert_snapshot!(run.snapshot, @" ClusteredFinalAggregateStream[0] register can_spill=true ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] @@ -737,20 +754,24 @@ async fn clustered_completed_groups() -> Result<()> { => Ok(1024 rows in 16 batches) "); - let run = run_final_and_single(|mode| clustered(mode, 34_000)).await?; + let run = run_final_and_single(|mode| clustered(mode, 38_500)).await?; assert_eq!(run.spill_count(), 0); assert_eq!(run.max_output_rows(), 256); - assert_snapshot!(run.snapshot, @r" + assert_snapshot!(run.snapshot, @" ClusteredFinalAggregateStream[0] register can_spill=true ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] @@ -768,20 +789,26 @@ async fn clustered_spills_then_replays() -> Result<()> { let _bound = hide_byte_counts().bind_to_scope(); let run = run_final_and_single(|mode| Case { batch_size: 256, - memory_limit: 34_000, + memory_limit: 40_000, ..case(mode, &["a", "b"], true, clustered_partial_input_long_run()) }) .await?; assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); assert_eq!(run.spill_count(), 2); - assert_snapshot!(run.snapshot, @r" + assert_snapshot!(run.snapshot, @" ClusteredFinalAggregateStream[0] register can_spill=true ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] shrink [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused @@ -790,15 +817,19 @@ async fn clustered_spills_then_replays() -> Result<()> { ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] + ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] From 339db8efb6328ba8b49f9d64035a43e74b863094 Mon Sep 17 00:00:00 2001 From: Yash Verma Date: Thu, 5 Nov 2026 17:27:09 +0530 Subject: [PATCH 2/4] test: leave headroom for final aggregate state batches --- .../sqllogictest/test_files/aggregate_memory_spill.slt | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index e8b9cdae130e8..6231cff37ddbf 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -226,9 +226,10 @@ FROM ( # --- Case G: partial/final aggregation under memory limit --- -# Four final aggregate partitions share the runtime pool; leave replay headroom. +# Four partial and final aggregate partitions share the runtime pool. Leave +# room to reserve incoming state batches during final merge and spill replay. statement ok -SET datafusion.runtime.memory_limit = '1M' +SET datafusion.runtime.memory_limit = '2M' statement ok SET datafusion.execution.target_partitions = 4 From a7850b71a8662ebc8261ce131a0769a13b2ecb3d Mon Sep 17 00:00:00 2001 From: Yash Verma Date: Thu, 5 Nov 2026 18:02:09 +0530 Subject: [PATCH 3/4] fix: allow bounded state batches under memory pressure --- .../src/aggregates/aggregate_stream.rs | 8 ++--- .../src/aggregates/clustered_final_stream.rs | 12 ++++--- .../src/aggregates/hash_stream.rs | 12 ++++--- .../physical-plan/src/aggregates/mod.rs | 29 ++++++++++------- .../cases/aggregate_memory_call_sequence.rs | 31 ------------------- .../test_files/aggregate_memory_spill.slt | 5 ++- 6 files changed, 39 insertions(+), 58 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 8cfa8457ce803..21a6e425463b8 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -13,6 +13,7 @@ use crate::aggregates::{ AccumulatorItem, AggrDynFilter, AggregateInputMode, AggregateMode, AggregateOutputMode, DynamicFilterAggregateType, aggregate_expressions, aggregate_metric_label, create_accumulators_with_metrics, + incoming_state_batch_reservation_bytes, }; use crate::metrics::{BaselineMetrics, RecordOutput}; use crate::stream::EmptyRecordBatchStream; @@ -20,7 +21,6 @@ use crate::{RecordBatchStream, SendableRecordBatchStream}; use arrow::array::ArrayRef; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{Result, ScalarValue, internal_datafusion_err, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::Operator; @@ -391,10 +391,10 @@ impl AggregateStream { // before merging it, then transfer that reservation to the // accumulator's state below. let previous_size = this.reservation.size(); + let batch_memory = incoming_state_batch_reservation_bytes(&batch); if this.mode.input_mode() == AggregateInputMode::Partial - && let Err(e) = this - .reservation - .try_grow(get_record_batch_memory_size(&batch)) + && batch_memory > 0 + && let Err(e) = this.reservation.try_grow(batch_memory) { this.finished = true; return Some((Err(e), this)); diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 26650857df324..38e0e10f6c8e4 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -9,18 +9,17 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::StreamExt; -use super::AggregateExec; use super::aggregate_hash_table::{ ClusteredAggregateTable, ClusteredAggregateTableMetrics, FinalMarker, }; use super::order::GroupClusteringMode; use super::spill::AggregateSpill; +use super::{AggregateExec, incoming_state_batch_reservation_bytes}; use crate::SendableRecordBatchStream; use crate::aggregates::AggregateMode; use crate::metrics::{BaselineMetrics, SpillMetrics}; @@ -324,8 +323,13 @@ impl Aggregating { while let Some(batch) = self.input.next().await.transpose()? { let timer = elapsed_compute.timer(); - let batch_memory = get_record_batch_memory_size(&batch); - match reservation.try_grow(batch_memory) { + let batch_memory = incoming_state_batch_reservation_bytes(&batch); + let admission = if batch_memory > 0 { + reservation.try_grow(batch_memory) + } else { + Ok(()) + }; + match admission { Ok(()) => {} Err(e @ DataFusionError::ResourcesExhausted(_)) => { let Some(spill_context) = self.spill_context.as_mut() else { diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 3e64d8e6956e5..dacd7d4b98e86 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -12,7 +12,6 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{ DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, }; @@ -20,7 +19,6 @@ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::{Stream, StreamExt}; -use super::AggregateExec; use super::aggregate_hash_table::{ AggregateHashTable, ClusteredAggregateTableMetrics, FinalMarker, PartialMarker, PartialSkipMarker, @@ -28,6 +26,7 @@ use super::aggregate_hash_table::{ use super::order::GroupClusteringMode; use super::skip_partial::SkipAggregationProbe; use super::spill::AggregateSpill; +use super::{AggregateExec, incoming_state_batch_reservation_bytes}; use crate::metrics::{ BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, }; @@ -683,8 +682,13 @@ impl FinalHashAggregateStream { let _timer = elapsed_compute.timer(); // Partial states may retain input-sized buffers. Account for the // incoming batch before merging it into the table. - let batch_memory = get_record_batch_memory_size(&batch); - match self.reservation.try_grow(batch_memory) { + let batch_memory = incoming_state_batch_reservation_bytes(&batch); + let admission = if batch_memory > 0 { + self.reservation.try_grow(batch_memory) + } else { + Ok(()) + }; + match admission { Ok(()) => {} Err(e @ DataFusionError::ResourcesExhausted(_)) => { // Spill the current table to make room for this batch, then diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 733e2fca656c6..1961f923acdbe 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -217,6 +217,20 @@ mod skip_partial; mod spill; mod topk; +/// Bounded batches passing between operators are not charged to the memory pool. +/// Partial-state batches larger than this may hold input-sized values and must +/// be reserved while the final aggregate merges them. +const MAX_IN_FLIGHT_STATE_BATCH_BYTES: usize = 64 * 1024; + +fn incoming_state_batch_reservation_bytes(batch: &RecordBatch) -> usize { + let bytes = datafusion_common::utils::memory::get_record_batch_memory_size(batch); + if bytes > MAX_IN_FLIGHT_STATE_BATCH_BYTES { + bytes + } else { + 0 + } +} + /// Hard-coded seed for aggregations to ensure hash values differ from `RepartitionExec`, avoiding collisions. const AGGREGATION_HASH_SEED: datafusion_common::hash_utils::RandomState = // This seed is chosen to be a large 64-bit number @@ -4275,18 +4289,9 @@ mod tests { assert!(final_stats.total_byte_size.get_value().is_some()); let task_ctx = if spill { - // The final table must be able to reserve one incoming partial-state - // batch as well as its initial footprint. Keep only a little extra - // room so growth of the table still triggers spilling. - let max_state_batch = result - .iter() - .map(datafusion_common::utils::memory::get_record_batch_memory_size) - .max() - .unwrap_or(0); - new_spill_ctx( - 2, - initial_reservation(&merged_aggregate)? + 2 * (max_state_batch + 64), - ) + // The fair pool is shared by the partial and final tables, give each + // one its initial table plus a little, so the final aggregation spills + new_spill_ctx(2, initial_reservation(&merged_aggregate)? + 2 * 64) } else { Arc::clone(&task_ctx) }; diff --git a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs index 6e2ea4dba710b..98e882b7a3092 100644 --- a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs +++ b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs @@ -511,16 +511,10 @@ async fn spills_twice_then_replays() -> Result<()> { FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] -> refused FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] -> refused FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] shrink [bytes] @@ -532,27 +526,21 @@ async fn spills_twice_then_replays() -> Result<()> { FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] shrink [bytes] @@ -616,7 +604,6 @@ async fn output_materialization_refused() -> Result<()> { FinalHashAggregateStream[0] register can_spill=true FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] -> refused FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] unregister @@ -737,15 +724,11 @@ async fn clustered_completed_groups() -> Result<()> { ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] @@ -762,16 +745,12 @@ async fn clustered_completed_groups() -> Result<()> { ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] @@ -800,15 +779,9 @@ async fn clustered_spills_then_replays() -> Result<()> { ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] shrink [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused @@ -817,19 +790,15 @@ async fn clustered_spills_then_replays() -> Result<()> { ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 6231cff37ddbf..e8b9cdae130e8 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -226,10 +226,9 @@ FROM ( # --- Case G: partial/final aggregation under memory limit --- -# Four partial and final aggregate partitions share the runtime pool. Leave -# room to reserve incoming state batches during final merge and spill replay. +# Four final aggregate partitions share the runtime pool; leave replay headroom. statement ok -SET datafusion.runtime.memory_limit = '2M' +SET datafusion.runtime.memory_limit = '1M' statement ok SET datafusion.execution.target_partitions = 4 From b1f2c298a8feabe6169fad8ddb0a8a13dd79c66e Mon Sep 17 00:00:00 2001 From: Yash Verma Date: Wed, 11 Nov 2026 17:07:40 +0530 Subject: [PATCH 4/4] fix: address reviewer feedback on aggregate state memory reservation - Charge partial-state batches only when bytes per row exceed 1 KiB and total bytes exceed 64 KiB - Restore early break on soft group limit before table memory resize in FinalHashAggregateStream - Extract shared spill_table helper in FinalHashAggregateStream - Guard incoming batch memory reservation calculation to Partial input mode in AggregateStream - Replace synthetic tests with large-buffer retaining batches to properly test state batch admission - Add test for ClusteredFinalAggregateStream spill-and-retry on partial clustering - Add test for FinalHashAggregateStream spill-and-retry on large state batch - Restore aggregate_memory_call_sequence.rs to match main --- .../src/aggregates/aggregate_stream.rs | 26 +-- .../src/aggregates/clustered_final_stream.rs | 99 +++++++++- .../src/aggregates/hash_stream.rs | 186 ++++++++++++------ .../physical-plan/src/aggregates/mod.rs | 22 ++- .../cases/aggregate_memory_call_sequence.rs | 26 +-- 5 files changed, 256 insertions(+), 103 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 21a6e425463b8..9b50670219d3b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -391,9 +391,13 @@ impl AggregateStream { // before merging it, then transfer that reservation to the // accumulator's state below. let previous_size = this.reservation.size(); - let batch_memory = incoming_state_batch_reservation_bytes(&batch); - if this.mode.input_mode() == AggregateInputMode::Partial - && batch_memory > 0 + let batch_memory = + if this.mode.input_mode() == AggregateInputMode::Partial { + incoming_state_batch_reservation_bytes(&batch) + } else { + 0 + }; + if batch_memory > 0 && let Err(e) = this.reservation.try_grow(batch_memory) { this.finished = true; @@ -608,7 +612,7 @@ mod tests { use crate::metrics::{MetricValue, MetricsSet}; use crate::test::TestMemoryExec; use crate::{ExecutionPlan, collect}; - use arrow::array::{Float64Array, Int32Array, Int64Array, UInt32Array}; + use arrow::array::{ArrayRef, Float64Array, Int32Array, Int64Array, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::DataFusionError; use datafusion_execution::memory_pool::GreedyMemoryPool; @@ -1059,13 +1063,11 @@ mod tests { empty, Arc::clone(&schema), )?; - // One batch may combine many partial states. The final SUM retains - // only one number, so reserving its accumulator alone misses the batch. + // 64 rows that retain 1 MB, like the few rows of a `median` state. let state_schema = partial.schema(); - let batch = RecordBatch::try_new( - Arc::clone(&state_schema), - vec![Arc::new(Int64Array::from(vec![1_i64; 16_384]))], - )?; + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64)); + let batch = RecordBatch::try_new(Arc::clone(&state_schema), vec![values])?; let input = TestMemoryExec::try_new_exec( &[vec![batch]], Arc::clone(&state_schema), @@ -1096,7 +1098,7 @@ mod tests { )); let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024 * 1024))) + .with_memory_pool(Arc::new(GreedyMemoryPool::new(2 * 1024 * 1024))) .build_arc()?; let context = Arc::new(TaskContext::default().with_runtime(runtime)); let output = collect(final_aggregate as Arc, context).await?; @@ -1105,7 +1107,7 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - assert_eq!(result.value(0), 16_384); + assert_eq!(result.value(0), 64); Ok(()) } } diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 38e0e10f6c8e4..9e0b14f548758 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -491,7 +491,7 @@ mod tests { }; use crate::common::collect; use crate::test::TestMemoryExec; - use arrow::array::{Int64Array, StringViewArray}; + use arrow::array::{ArrayRef, Int64Array, StringViewArray}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; @@ -790,16 +790,99 @@ mod tests { assert_resources_exhausted_under_limit(&plan, limit).await } + /// The table stays small: only the incoming batch can exceed the limit. #[tokio::test] - async fn clustered_final_reserves_incoming_state_batch() -> Result<()> { - // Many partial SUM states for one key merge to a small final state. - let plan = sum_over_sorted_input( + async fn clustered_final_refuses_large_state_batch() -> Result<()> { + let plan = + sum_over_sorted_input(AggregateMode::Final, &["k"], vec![vec![vec![1; 64]]])?; + // 64 rows that retain 1 MB per column, like the few rows of a `median` state. + let state = plan.input().schema(); + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64)); + let batch = + RecordBatch::try_new(Arc::clone(&state), vec![Arc::clone(&values), values])?; + let ordering = plan.input().properties().output_ordering().unwrap().clone(); + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&state), None)? + .try_with_sort_information(vec![ordering])?; + let plan = AggregateExec::try_new( AggregateMode::Final, - &["k"], - vec![vec![vec![1; 16_384]]], + plan.group_expr().clone(), + plan.aggr_expr().to_vec(), + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + plan.input_schema(), )?; assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); - assert_resources_exhausted_under_limit(&plan, 64 * 1024).await?; - assert_output_unchanged_under_limit(&plan, 1024 * 1024).await + assert_resources_exhausted_under_limit(&plan, 1024 * 1024).await?; + assert_output_unchanged_under_limit(&plan, 4 * 1024 * 1024).await + } + + /// Spilling the table makes room for a state batch that did not fit next to it. + #[tokio::test] + async fn clustered_final_spills_to_admit_large_state_batch() -> Result<()> { + let plan = sum_over_sorted_input( + AggregateMode::Final, + &["k1", "k2"], + vec![vec![vec![0; 4000], (0..4000).collect()]], + )?; + assert_eq!( + plan.group_clustering_mode(), + &GroupClusteringMode::Partial(vec![0]) + ); + let state = plan.input().schema(); + let batch = |rows: i64| { + RecordBatch::try_new( + Arc::clone(&state), + vec![ + Arc::new(Int64Array::from(vec![0_i64; rows as usize])) as ArrayRef, + Arc::new(Int64Array::from((0..rows).collect::>())), + Arc::new(Int64Array::from(vec![1_i64; rows as usize])), + ], + ) + }; + // 64 rows that retain 3 MB, like the few rows of a `median` state. + let large = batch(131_072)?.slice(0, 64); + let ordering = plan.input().properties().output_ordering().unwrap().clone(); + let input = TestMemoryExec::try_new( + &[vec![batch(4000)?, large]], + Arc::clone(&state), + None, + )? + .try_with_sort_information(vec![ordering])?; + let plan = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + plan.group_expr().clone(), + plan.aggr_expr().to_vec(), + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + plan.input_schema(), + )?); + + // Fits the large batch, but not next to the table of 4000 groups. + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(3150 * 1024))) + .build_arc()?; + let context = + Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + )) + .with_runtime(runtime), + ); + let output = collect(plan.execute(0, context)?).await?; + assert!(plan.metrics().unwrap().spill_count().unwrap() > 0); + let rows: usize = output.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 4000); + let mut total = 0; + for b in &output { + total += arrow::compute::sum( + b.column(2).as_any().downcast_ref::().unwrap(), + ) + .unwrap_or(0); + } + assert_eq!(total, 4064); + Ok(()) } } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index dacd7d4b98e86..7aadd61127d9d 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -12,9 +12,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{ - DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, -}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::{Stream, StreamExt}; @@ -667,6 +665,31 @@ impl FinalHashAggregateStream { } } + fn spill_table( + hash_table: &mut AggregateHashTable, + spill_context: Option<&mut Box>, + reservation: &MemoryReservation, + err: DataFusionError, + ) -> Result<()> { + let Some(spill_context) = spill_context else { + return Err(err.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + if hash_table.building_group_count() == 0 { + return Err(err); + } + let result = hash_table + .take_state_batch() + .and_then(|state| spill_context.sort_and_spill(state)); + reservation + .try_resize(hash_table.memory_size()) + .map_err(|e| { + e.context("Decreasing allocation after spilling should succeed") + })?; + result + } + /// Read input stream, if no memory, then spill and continue reading - aggregate partial state batches into the hash table. /// /// Spilling: The table cannot reserve enough memory. @@ -694,19 +717,12 @@ impl FinalHashAggregateStream { // Spill the current table to make room for this batch, then // retry admission. An individual batch that still cannot fit // must fail rather than proceed without a reservation. - let Some(spill_context) = spill_context.as_mut() else { - return Err(e.context( - "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", - )); - }; - if hash_table.building_group_count() == 0 { - return Err(e); - } - let result = hash_table - .take_state_batch() - .and_then(|state| spill_context.sort_and_spill(state)); - self.reservation.try_resize(hash_table.memory_size())?; - result?; + Self::spill_table( + hash_table, + spill_context.as_mut(), + &self.reservation, + e, + )?; self.reservation.try_grow(batch_memory)?; } Err(e) => return Err(e), @@ -719,7 +735,9 @@ impl FinalHashAggregateStream { let spilled = spill_context .as_ref() .is_some_and(|context| context.has_spills()); - let hit_soft_limit = self.hit_soft_group_limit(hash_table) && !spilled; + if self.hit_soft_group_limit(hash_table) && !spilled { + break; + } // Replace the input-batch reservation with the table footprint, // spilling if the table cannot be reserved. @@ -731,48 +749,17 @@ impl FinalHashAggregateStream { )); match resize_result { - Ok(()) => { - if hit_soft_limit { - break; - } - } + Ok(()) => {} // The table cannot reserve enough memory. // Move all current states into one fully group-key-sorted spill run. Err(e @ DataFusionError::ResourcesExhausted(_)) => { - // OOM and don't support spilling from configuration - let spill_context = spill_context.as_mut().ok_or_else(|| e.context( - "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", - ))?; - - // Sanity check: impossible to OOM when there is no group aggregated. - assert_ne_or_internal_err!( - hash_table.building_group_count(), - 0, - "Final hash aggregate ran out of memory with no aggregated groups" - ); - - // Sorts and spills one complete in-memory state run - - // Go to the next state to perform spilling the aggregated - // groups so far. - let result = hash_table - .take_state_batch() - .and_then(|batch| spill_context.sort_and_spill(batch)); - - // Spilling shrinks the aggregate table and releases its accumulated - // memory. Update the reservation accordingly. - self.reservation - .try_resize(hash_table.memory_size()) - .map_err(|e| { - e.context( - "Decreasing allocation after spilling should succeed", - ) - })?; - - result?; - - // One sorted run was written; resume reading the original input. + Self::spill_table( + hash_table, + spill_context.as_mut(), + &self.reservation, + e, + )?; } Err(e) => return Err(e), } @@ -1794,12 +1781,12 @@ mod tests { Arc::clone(&schema), )?; let state_schema = partial.schema(); + // 64 rows that retain 2 MB, like the few rows of a `median` state. + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64)); let batch = RecordBatch::try_new( Arc::clone(&state_schema), - vec![ - Arc::new(Int64Array::from(vec![1_i64; 16_384])), - Arc::new(Int64Array::from(vec![1_i64; 16_384])), - ], + vec![Arc::clone(&values), values], )?; let input = TestMemoryExec::try_new_exec( &[vec![batch]], @@ -1838,7 +1825,7 @@ mod tests { )); let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024 * 1024))) + .with_memory_pool(Arc::new(GreedyMemoryPool::new(3 * 1024 * 1024))) .build_arc()?; let context = Arc::new( TaskContext::default() @@ -1851,7 +1838,84 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - assert_eq!(result.value(0), 16_384); + assert_eq!(result.value(0), 64); + Ok(()) + } + + /// Spilling the table makes room for a state batch that did not fit next to it. + #[tokio::test] + async fn final_hash_spills_to_admit_large_state_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]); + let aggregate = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + ); + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let state = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![None], + empty, + Arc::clone(&schema), + )? + .schema(); + let batch = |rows: i64| { + RecordBatch::try_new( + Arc::clone(&state), + vec![ + Arc::new(Int64Array::from((0..rows).collect::>())) as ArrayRef, + Arc::new(Int64Array::from(vec![1_i64; rows as usize])), + ], + ) + }; + // 64 rows that retain 2 MB, like the few rows of a `median` state. + let large = batch(131_072)?.slice(0, 64); + let input = TestMemoryExec::try_new_exec( + &[vec![batch(4000)?, large]], + Arc::clone(&state), + None, + )?; + let plan = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggregate], + vec![None], + input, + Arc::clone(&state), + )?); + // Fits the large batch, but not next to the table of 4000 groups. + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(2200 * 1024))) + .build_arc()?; + let context = + Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + )) + .with_runtime(runtime), + ); + let output = collect(plan.execute(0, context)?).await?; + assert!(plan.metrics().unwrap().spill_count().unwrap() > 0); + let rows: usize = output.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 4000); + let mut total = 0; + for b in &output { + total += arrow::compute::sum( + b.column(1).as_primitive::(), + ) + .unwrap_or(0); + } + assert_eq!(total, 4064); Ok(()) } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 1961f923acdbe..471da9f11acc7 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -217,18 +217,22 @@ mod skip_partial; mod spill; mod topk; -/// Bounded batches passing between operators are not charged to the memory pool. -/// Partial-state batches larger than this may hold input-sized values and must -/// be reserved while the final aggregate merges them. -const MAX_IN_FLIGHT_STATE_BATCH_BYTES: usize = 64 * 1024; +/// Batches passing between operators are bounded by `batch_size` rows and are +/// not charged to the memory pool. A partial-state batch is charged only when +/// its rows carry more than this many bytes on average, which means a row holds +/// many input values (`median`, `array_agg`), so its size follows the input. +const MAX_IN_FLIGHT_STATE_BYTES_PER_ROW: usize = 1024; + +/// Below this size a batch is never charged, whatever its row count. +const MIN_CHARGED_STATE_BATCH_BYTES: usize = 64 * 1024; fn incoming_state_batch_reservation_bytes(batch: &RecordBatch) -> usize { let bytes = datafusion_common::utils::memory::get_record_batch_memory_size(batch); - if bytes > MAX_IN_FLIGHT_STATE_BATCH_BYTES { - bytes - } else { - 0 - } + let in_flight_bytes = batch + .num_rows() + .saturating_mul(MAX_IN_FLIGHT_STATE_BYTES_PER_ROW) + .max(MIN_CHARGED_STATE_BATCH_BYTES); + if bytes > in_flight_bytes { bytes } else { 0 } } /// Hard-coded seed for aggregations to ensure hash values differ from `RepartitionExec`, avoiding collisions. diff --git a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs index 98e882b7a3092..853cbabf3ce50 100644 --- a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs +++ b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs @@ -446,16 +446,16 @@ fn hide_byte_counts() -> Settings { /// Runs `case` as Final and as Single and returns the Final run. /// -/// Final reserves each incoming partial-state batch before merging it, while -/// Single consumes raw rows. Their pool calls differ, but their results must -/// still match. The test snapshots the Final run. +/// The Final and Single streams account in the same way, so this asserts that +/// both make the same pool calls, with the same byte counts, once the stream +/// name is replaced, and the test snapshots only the Final run. async fn run_final_and_single(case: impl Fn(AggregateMode) -> Case) -> Result { let final_run = case(AggregateMode::Final).run().await?; let single_run = case(AggregateMode::Single).run().await?; assert_eq!( - datafusion_common::test_util::batches_to_sort_string(&single_run.output), - datafusion_common::test_util::batches_to_sort_string(&final_run.output), - "Final and Single return different results" + single_run.snapshot.replace("Single", "Final"), + final_run.snapshot, + "Final and Single make different pool calls" ); Ok(final_run) } @@ -506,7 +506,7 @@ async fn spills_twice_then_replays() -> Result<()> { .await?; assert_eq!(run.clustering, GroupClusteringMode::None); assert_eq!(run.spill_count(), 2); - assert_snapshot!(run.snapshot, @" + assert_snapshot!(run.snapshot, @r" FinalHashAggregateStream[0] register can_spill=true FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] @@ -600,7 +600,7 @@ async fn output_materialization_refused() -> Result<()> { let run = padded(AggregateMode::Final).run().await?; assert_eq!(run.spill_count(), 1); assert_eq!(run.output_batch_rows, vec![64; 4]); - assert_snapshot!(run.snapshot, @" + assert_snapshot!(run.snapshot, @r" FinalHashAggregateStream[0] register can_spill=true FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] @@ -718,7 +718,7 @@ async fn clustered_completed_groups() -> Result<()> { assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); assert_eq!(run.spill_count(), 0); assert_eq!(run.max_output_rows(), 64); - assert_snapshot!(run.snapshot, @" + assert_snapshot!(run.snapshot, @r" ClusteredFinalAggregateStream[0] register can_spill=true ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] @@ -737,10 +737,10 @@ async fn clustered_completed_groups() -> Result<()> { => Ok(1024 rows in 16 batches) "); - let run = run_final_and_single(|mode| clustered(mode, 38_500)).await?; + let run = run_final_and_single(|mode| clustered(mode, 34_000)).await?; assert_eq!(run.spill_count(), 0); assert_eq!(run.max_output_rows(), 256); - assert_snapshot!(run.snapshot, @" + assert_snapshot!(run.snapshot, @r" ClusteredFinalAggregateStream[0] register can_spill=true ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] @@ -768,13 +768,13 @@ async fn clustered_spills_then_replays() -> Result<()> { let _bound = hide_byte_counts().bind_to_scope(); let run = run_final_and_single(|mode| Case { batch_size: 256, - memory_limit: 40_000, + memory_limit: 34_000, ..case(mode, &["a", "b"], true, clustered_partial_input_long_run()) }) .await?; assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); assert_eq!(run.spill_count(), 2); - assert_snapshot!(run.snapshot, @" + assert_snapshot!(run.snapshot, @r" ClusteredFinalAggregateStream[0] register can_spill=true ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes]