From c95401e108e052e2690955d68ef74f43c49bc4ff Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 21 Oct 2026 10:45:02 -0700 Subject: [PATCH 1/2] fix: emit completed groups when a fully sorted partial aggregate runs out of memory --- .../aggregates/clustered_partial_stream.rs | 108 ++++++++++++++++-- 1 file changed, 101 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index f7f7e2d9031dd..5fab812da1b92 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -85,8 +85,8 @@ use crate::{SendableRecordBatchStream, metrics}; /// 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, the stream still emits the +/// completed groups, and returns the error only when no group is complete. /// /// ## Partial group clustering /// @@ -318,12 +318,19 @@ impl Aggregating { let output = match reservation.try_resize(self.table.memory_size()) { Ok(()) => self.table.take_completed_state_batch()?, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { - // Partial group clustering may have an unbounded active key range. + // Emitting the completed groups leaves only the active group + // under full group clustering. Partial group clustering may + // have an unbounded active key range, so emit every state. // The final stage can merge incomplete states emitted here. - if matches!(self.table.group_clustering(), GroupClustering::Full(_)) { - return Err(oom); - } - let Some(batch) = self.table.take_state_batch()? else { + let output = if matches!( + self.table.group_clustering(), + GroupClustering::Full(_) + ) { + self.table.take_completed_state_batch()? + } else { + self.table.take_state_batch()? + }; + let Some(batch) = output else { return Err(oom); }; Some(batch) @@ -419,3 +426,90 @@ impl Outputting { Ok(next_stage) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::aggregates::PhysicalGroupBy; + use crate::test::TestMemoryExec; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::test_util::batches_to_string; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, + }; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::LexOrdering; + + /// Runs a partial `SUM(v) GROUP BY k` over three batches sorted on `k`, where + /// each batch boundary repeats a key. Returns the output and the most memory + /// the stream held between batches. + async fn run(pool: Arc) -> Result<(Vec, usize)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batches = [0..=256, 256..=512, 512..=768] + .into_iter() + .map(|keys| { + let k = Int64Array::from_iter_values(keys); + let v = Int64Array::from(vec![1; k.len()]); + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(k), Arc::new(v)]) + }) + .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::Partial, + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + )], + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + Arc::clone(&schema), + )?; + assert_eq!( + 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 = + ClusteredPartialAggregateStream::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(Arc::new(UnboundedMemoryPool::default())).await?; + let (actual, _) = run(Arc::new(GreedyMemoryPool::new(held))).await?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + Ok(()) + } +} From 2ad2ab4db6af6548a13e108ab0d47068c9bd58fc Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 22 Oct 2026 17:42:21 -0700 Subject: [PATCH 2/2] Simplify the OOM arm and the memory-pressure test --- .../aggregates/clustered_partial_stream.rs | 48 +++++++------------ 1 file changed, 16 insertions(+), 32 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index 5fab812da1b92..f60f7ed30cdc7 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -85,8 +85,9 @@ use crate::{SendableRecordBatchStream, metrics}; /// only the current group remains active after completed groups are emitted, and /// memory usage does not grow with the total number of groups. /// -/// If the reservation after an input batch fails, the stream still emits the -/// completed groups, and returns the error only when no group is complete. +/// If the reservation after an input batch fails, `ClusteredPartialAggregateStream` +/// still emits the completed groups, and returns the error only when no group is +/// complete. /// /// ## Partial group clustering /// @@ -318,22 +319,13 @@ impl Aggregating { let output = match reservation.try_resize(self.table.memory_size()) { Ok(()) => self.table.take_completed_state_batch()?, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { - // Emitting the completed groups leaves only the active group - // under full group clustering. Partial group clustering may - // have an unbounded active key range, so emit every state. + // Partial group clustering may have an unbounded active key range. // The final stage can merge incomplete states emitted here. - let output = if matches!( - self.table.group_clustering(), - GroupClustering::Full(_) - ) { - self.table.take_completed_state_batch()? + if matches!(self.table.group_clustering(), GroupClustering::Full(_)) { + Some(self.table.take_completed_state_batch()?.ok_or(oom)?) } else { - self.table.take_state_batch()? - }; - let Some(batch) = output else { - return Err(oom); - }; - Some(batch) + Some(self.table.take_state_batch()?.ok_or(oom)?) + } } Err(e) => return Err(e), }; @@ -445,20 +437,16 @@ mod tests { use datafusion_physical_expr::expressions::col; use datafusion_physical_expr_common::sort_expr::LexOrdering; - /// Runs a partial `SUM(v) GROUP BY k` over three batches sorted on `k`, where - /// each batch boundary repeats a key. Returns the output and the most memory - /// the stream held between batches. + /// Runs a partial `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 between batches. async fn run(pool: Arc) -> Result<(Vec, usize)> { - let schema = Arc::new(Schema::new(vec![ - Field::new("k", DataType::Int64, false), - Field::new("v", DataType::Int64, false), - ])); - let batches = [0..=256, 256..=512, 512..=768] + 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); - let v = Int64Array::from(vec![1; k.len()]); - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(k), Arc::new(v)]) + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(k)]) }) .collect::, _>>()?; let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? @@ -470,19 +458,15 @@ mod tests { AggregateMode::Partial, PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]), vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + 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::new(input), Arc::clone(&schema), )?; - assert_eq!( - aggregate.group_clustering_mode(), - &GroupClusteringMode::Full - ); let context = Arc::new( TaskContext::default().with_runtime( RuntimeEnvBuilder::new()