From 10b67d482c0b372ca05dddae162ec979dd9a4cca Mon Sep 17 00:00:00 2001 From: Ran Reichman Date: Wed, 26 Aug 2026 07:59:25 -0400 Subject: [PATCH] charge accumulators' construction-time state to the pool in ungrouped aggregation --- .../src/aggregates/aggregate_stream.rs | 5 +++++ .../physical-plan/src/aggregates/mod.rs | 19 +++++++++++++++---- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index ac7727b459300..f9f0ff3dc6acb 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -301,6 +301,11 @@ impl AggregateStream { let reservation = MemoryConsumer::new(format!("AggregateStream[{partition}]")) .register(context.memory_pool()); + // Accumulators that allocate their state at construction are invisible to the + // per-batch size deltas in `aggregate_batch`, so charge their initial sizes up + // front (mirroring `GroupsAccumulatorAdapter`, which charges `state.size()` when + // it creates each accumulator). + reservation.try_grow(accumulators.iter().map(|accum| accum.size()).sum())?; // Enable dynamic filter if: // 1. AggregateExec did the check and ensure it supports the dynamic filter diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 2dc41adfd6c62..3499ab988c870 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4011,21 +4011,32 @@ mod tests { Arc::clone(&input_schema), )?); - let stream = partial_aggregate.execute_typed(0, &task_ctx)?; + let stream = partial_aggregate.execute_typed(0, &task_ctx); // ensure that we really got the version we wanted - match version { + let stream = match version { 0 => { - assert!(matches!(stream, StreamType::AggregateStream(_))); + // the ungrouped stream charges its accumulators' construction-time + // state up front, so admission fails before any input is read + let err = stream.err().unwrap(); + assert!( + matches!(err.find_root(), DataFusionError::ResourcesExhausted(_)), + "Wrong error type: {err}", + ); + continue; } 1 => { + let stream = stream?; assert!(matches!(stream, StreamType::GroupedHash(_))); + stream } 2 => { + let stream = stream?; assert!(matches!(stream, StreamType::SingleHash(_))); + stream } _ => panic!("Unknown version: {version}"), - } + }; let stream: SendableRecordBatchStream = stream.into(); let err = collect(stream).await.unwrap_err();