From 1ca42c8e90e1667d64d49eace9f0310015c081c0 Mon Sep 17 00:00:00 2001 From: yashaggarwal85d <68061195+yashaggarwal85d@users.noreply.github.com> Date: Fri, 23 Oct 2026 04:55:18 +0000 Subject: [PATCH] fix: emit completed groups when fully sorted final and single aggregates run out of memory When a memory reservation is refused and the stream cannot spill, `ClusteredFinalAggregateStream` and `ClusteredSingleAggregateStream` now emit the completed groups and keep only the active group, returning the error only when no group is complete. This mirrors the fix for the partial stream. Closes #26628. --- .../src/aggregates/clustered_final_stream.rs | 101 +++++++++++++++++- .../src/aggregates/clustered_single_stream.rs | 97 ++++++++++++++++- 2 files changed, 190 insertions(+), 8 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 9c71c5ee1332..ed5cd839eced 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -261,7 +261,8 @@ impl ClusteredFinalAggregateStream { /// 1. Start. /// 2. Merge one input batch: /// - If memory fits and no groups are complete, continue reading input. - /// - If OOM, spill. + /// - If OOM, spill. Without spilling, emit completed groups through [3], + /// and fail only when no group is complete. /// 3. Prepare output: /// - Before any spill, a prefix of groups is complete: materialize the /// entire prefix once, retaining the input and active groups to resume @@ -340,7 +341,15 @@ impl Aggregating { Ok(()) => {} Err(oom @ DataFusionError::ResourcesExhausted(_)) => { let Some(spill_context) = self.spill_context.as_mut() else { - return Err(oom); + // Without spilling, emitting completed groups is the only + // way to release memory. Fail only when none is complete. + let batch = + self.table.take_completed_result_batch()?.ok_or(oom)?; + timer.done(); + return Ok(Some(ExecutionStage::Outputting(Outputting { + batch, + resume: Some(self), + }))); }; if self.table.is_empty() { return Err(oom); @@ -464,8 +473,11 @@ mod tests { use crate::test::TestMemoryExec; use arrow::array::{Int64Array, StringViewArray}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::test_util::batches_to_string; use datafusion_execution::config::SessionConfig; - use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, + }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::{min_max::min_udaf, sum::sum_udaf}; use datafusion_physical_expr::PhysicalSortExpr; @@ -732,4 +744,87 @@ mod tests { assert_eq!(pool.reserved(), 0); Ok(()) } + + /// Runs a final `SUM(k) GROUP BY k` over two partial-state batches sorted on + /// `k` that share a boundary key. Returns the output and the most memory the + /// stream held while handing off output. + async fn run_fully_sorted( + pool: Arc, + ) -> Result<(Vec, usize)> { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let groups = PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]); + let expressions = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("k", &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, + groups.clone(), + expressions.clone(), + vec![None], + empty, + Arc::clone(&schema), + )?; + let partial_schema = partial.schema(); + let batches = [0..=256, 256..=512] + .into_iter() + .map(|keys| { + let k = Int64Array::from_iter_values(keys); + RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![Arc::new(k.clone()), Arc::new(k)], + ) + }) + .collect::, _>>()?; + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("k", &partial_schema)?)]) + .unwrap(); + let input = + TestMemoryExec::try_new(&[batches], Arc::clone(&partial_schema), None)? + .try_with_sort_information(vec![ordering])?; + let aggregate = AggregateExec::try_new( + AggregateMode::Final, + groups.as_final(), + expressions, + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + schema, + )?; + assert!(matches!( + aggregate.group_clustering_mode(), + GroupClusteringMode::Full + )); + let context = Arc::new( + TaskContext::default().with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?, + ), + ); + let mut stream = + ClusteredFinalAggregateStream::new(&aggregate, &context, 0)?.into_stream(); + // Each output batch is a single slice, so the stream releases it before + // handing it off and holds only the table. + let mut held = pool.reserved(); + let mut output = vec![]; + while let Some(batch) = stream.next().await.transpose()? { + held = held.max(pool.reserved()); + output.push(batch); + } + Ok((output, held)) + } + + /// The limit fits the table with only the active group, but not the table + /// before the completed groups are emitted. The result must not change. + #[tokio::test] + async fn completed_groups_under_memory_pressure() -> Result<()> { + let (expected, held) = + run_fully_sorted(Arc::new(UnboundedMemoryPool::default())).await?; + let (actual, _) = run_fully_sorted(Arc::new(GreedyMemoryPool::new(held))).await?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs index 7095921ac78f..7f47eb8e6695 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs @@ -82,8 +82,9 @@ use crate::{RecordBatchStream, SendableRecordBatchStream}; /// only the current group remains active after completed groups are emitted, and /// memory usage does not grow with the total number of groups. /// -/// If a memory reservation nevertheless fails, the stream returns the error -/// directly, indicating an unexpected behavior. +/// If the reservation after an input batch fails, `ClusteredSingleAggregateStream` +/// still emits the completed groups, and returns the error only when no group is +/// complete. /// /// ## Partial group clustering /// @@ -335,7 +336,27 @@ impl ClusteredSingleAggregateStream { let Some(spill_context) = spill_context else { // `None` means spilling is not supported, see comments // at `ClusteredSingleAggregateState` for details. - return Self::break_with_err(e); + // Emitting completed groups is then the only way to + // release memory. Fail only when none is complete. + let timer = elapsed_compute.timer(); + let result = table.take_completed_result_batch(); + timer.done(); + return match result { + Ok(Some(batch)) => { + let table_memory = + Self::reservation_size_for_table(&table, None); + self.start_outputting( + batch, + table_memory, + ClusteredSingleAggregateState::ReadingInput { + table, + spill_context: None, + }, + ) + } + Ok(None) => Self::break_with_err(e), + Err(e) => Self::break_with_err(e), + }; }; if table.is_empty() { return Self::break_with_internal_err( @@ -766,10 +787,10 @@ mod tests { use crate::test::TestMemoryExec; use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::test_util::batches_to_sort_string; + use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{ - GreedyMemoryPool, MemoryPool, PeakRecordingPool, + GreedyMemoryPool, MemoryPool, PeakRecordingPool, UnboundedMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::sum::sum_udaf; @@ -851,4 +872,70 @@ mod tests { ); Ok(()) } + + /// Runs a single `SUM(k) GROUP BY k` over two batches sorted on `k` that + /// share a boundary key. Returns the output and the most memory the stream + /// held while handing off output. + async fn run_fully_sorted( + pool: Arc, + ) -> Result<(Vec, usize)> { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let batches = [0..=256, 256..=512] + .into_iter() + .map(|keys| { + let k = Int64Array::from_iter_values(keys); + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(k)]) + }) + .collect::, _>>()?; + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ + LexOrdering::new([PhysicalSortExpr::new_default(col("k", &schema)?)]) + .unwrap(), + ])?; + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("k", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + )], + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.group_clustering_mode, + GroupClusteringMode::Full + )); + let context = Arc::new( + TaskContext::default().with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?, + ), + ); + let mut stream = ClusteredSingleAggregateStream::new(&aggregate, &context, 0)?; + // Each output batch is a single slice, so the stream releases it before + // handing it off and holds only the table. + let mut held = pool.reserved(); + let mut output = vec![]; + while let Some(batch) = stream.next().await.transpose()? { + held = held.max(pool.reserved()); + output.push(batch); + } + Ok((output, held)) + } + + /// The limit fits the table with only the active group, but not the table + /// before the completed groups are emitted. The result must not change. + #[tokio::test] + async fn fully_sorted_completed_groups_under_memory_pressure() -> Result<()> { + let (expected, held) = + run_fully_sorted(Arc::new(UnboundedMemoryPool::default())).await?; + let (actual, _) = run_fully_sorted(Arc::new(GreedyMemoryPool::new(held))).await?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + Ok(()) + } }