From 86a9abbdc628bf948a99e729a8f5ba91b283d0ad Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 26 Oct 2026 17:00:26 -0700 Subject: [PATCH 1/2] fix: release per-batch accumulator scratch after a clustered partial emit Add GroupsAccumulator::shrink_to_fit, a no-op by default, and call it for every accumulator from ClusteredAggregateTable::shrink_to_fit. HllGroupsAccumulator drops its per-batch hashes, and GroupsAccumulatorAdapter drops its per-batch scratch, which it otherwise keeps while the active group remains. Closes #26781. Co-authored-by: Isaac --- .../expr-common/src/groups_accumulator.rs | 8 ++ .../src/aggregate/groups_accumulator.rs | 8 ++ .../src/approx_distinct.rs | 5 + .../aggregates/aggregate_hash_table/common.rs | 5 + .../aggregate_hash_table/common_clustered.rs | 10 +- .../aggregates/clustered_partial_stream.rs | 93 +++++++++++++++++-- 6 files changed, 117 insertions(+), 12 deletions(-) diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index e7301f85e068b..83d9ee2367e7d 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -368,6 +368,14 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; + + /// Keeps the stored groups but releases capacity they do not need, such as + /// per-batch scratch buffers, so that [`Self::size`] reflects the stored + /// groups. For example, after [`EmitTo::First`] removes most groups, an + /// accumulator may keep a buffer it sized for the last input batch. + /// + /// The default does nothing. + fn shrink_to_fit(&mut self) {} } #[cfg(test)] diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 8c86d59e44f37..a99448e3045d1 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -653,6 +653,14 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { self.allocation_bytes } + fn shrink_to_fit(&mut self) { + let vec_size_pre = self.states.allocated_size(); + self.states.shrink_to_fit(); + self.adjust_allocation(vec_size_pre, self.states.allocated_size()); + let scratch = std::mem::take(&mut self.scratch); + self.free_allocation(scratch.allocated_size()); + } + fn convert_to_state( &self, values: &[ArrayRef], diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index a925a994c33be..f9b9affc0e730 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -629,6 +629,11 @@ impl GroupsAccumulator for HllGroupsAccumulator { + self.allocated_bytes + self.hashes.capacity() * size_of::() } + + fn shrink_to_fit(&mut self) { + self.groups.shrink_to_fit(); + self.hashes = Vec::new(); + } } impl Debug for ApproxDistinct { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 12a7e9549172a..9a8919da13074 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -816,6 +816,11 @@ impl HashAggregateAccumulator { self.accumulator.size() } + /// See [`GroupsAccumulator::shrink_to_fit`]. + pub(super) fn shrink_to_fit(&mut self) { + self.accumulator.shrink_to_fit(); + } + pub(super) fn update_batch( &mut self, values: &CompactedAccumulatorArgs, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs index e75ed3012e70c..b98b7c456d55a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs @@ -287,14 +287,16 @@ impl ClusteredAggregateTable { } /// Releases the capacity that the buffer holds beyond its stored groups, - /// such as hash table slots and the per-batch `group_indices` scratch, so - /// that [`Self::memory_size`] reflects the remaining groups. - /// - /// Accumulator capacity is not released. + /// such as hash table slots, the per-batch `group_indices` scratch and + /// per-batch accumulator scratch, so that [`Self::memory_size`] reflects + /// the remaining groups. pub(in crate::aggregates) fn shrink_to_fit(&mut self) { self.buffer.group_values.shrink_to_fit(); self.buffer.group_indices.clear(); self.buffer.group_indices.shrink_to_fit(); + for acc in &mut self.buffer.accumulators { + acc.shrink_to_fit(); + } } pub(in crate::aggregates) fn metrics(&self) -> ClusteredAggregateTableMetrics { diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index 0232a6df8c6dc..886a80484e775 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -434,9 +434,14 @@ mod tests { GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::AggregateUDF; + use datafusion_functions_aggregate::any_value::any_value_udaf; + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::aggregate::{ + AggregateExprBuilder, AggregateFunctionExpr, + }; use datafusion_physical_expr::expressions::col; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -484,6 +489,21 @@ mod tests { async fn run( batches: Vec, pool: Arc, + ) -> Result> { + let schema = batches[0].schema(); + let sum = AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?; + run_aggregates(batches, vec![Arc::new(sum)], pool).await + } + + /// Runs a partial aggregate of `aggregates` grouped by `k0, ..` over + /// `batches`, which are sorted on every key. + async fn run_aggregates( + batches: Vec, + aggregates: Vec>, + pool: Arc, ) -> Result> { let schema = batches[0].schema(); let keys = (0..schema.fields().len() - 1) @@ -499,13 +519,8 @@ mod tests { let aggregate = AggregateExec::try_new( AggregateMode::Partial, PhysicalGroupBy::new_single(keys), - vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("sum") - .build()?, - )], - vec![None], + aggregates.clone(), + vec![None; aggregates.len()], Arc::new(input), Arc::clone(&schema), )?; @@ -559,4 +574,66 @@ mod tests { assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); assert_completed_groups_emitted_under_limit(&plan).await } + + /// Like [`completed_groups_under_memory_pressure`], but with accumulators + /// that keep per-batch scratch. For example, each `APPROX_DISTINCT` keeps + /// the hashes of the last batch (64 KB for 8192 rows), and the + /// `GroupsAccumulatorAdapter` of `ANY_VALUE` keeps its per-batch indices + /// (~160 KB). After the completed groups are emitted, the accumulators must + /// release that scratch too, or reserving the table fails the limit + /// although the active group needs a few KB. The output must not change. + #[tokio::test] + async fn accumulator_scratch_under_memory_pressure() -> Result<()> { + type Aggregates = fn(&SchemaRef) -> Result>>; + fn build( + udaf: Arc, + args: &[&str], + schema: &SchemaRef, + ) -> Result> { + let alias = format!("{}({})", udaf.name(), args.join(", ")); + let args = args + .iter() + .map(|name| col(name, schema)) + .collect::>>()?; + Ok(Arc::new( + AggregateExprBuilder::new(udaf, args) + .schema(Arc::clone(schema)) + .alias(alias) + .build()?, + )) + } + let cases: [(&str, Aggregates); 2] = [ + // `HllGroupsAccumulator` + ("approx_distinct", |schema| { + Ok(vec![ + build(approx_distinct_udaf(), &["k0"], schema)?, + build(approx_distinct_udaf(), &["v"], schema)?, + ]) + }), + // `GroupsAccumulatorAdapter` + ("any_value", |schema| { + Ok(vec![build(any_value_udaf(), &["v"], schema)?]) + }), + ]; + for (name, aggregates) in cases { + let input = + || batches(&[DataType::Int64], &[0..8193, 8192..16385, 16384..24576]); + let schema = input()?[0].schema(); + let expected = run_aggregates( + input()?, + aggregates(&schema)?, + Arc::new(UnboundedMemoryPool::default()), + ) + .await?; + let actual = run_aggregates( + input()?, + aggregates(&schema)?, + Arc::new(GreedyMemoryPool::new(100 * 1024)), + ) + .await + .map_err(|e| e.context(format!("aggregates {name}")))?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + } + Ok(()) + } } From f2f305d35071e67ddd731d5029e7c4a0993c02bb Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 28 Oct 2026 16:36:34 -0700 Subject: [PATCH 2/2] test: cover accumulator scratch release in clustered single aggregation The clustered final and single streams also call ClusteredAggregateTable::shrink_to_fit after the memory-pressure emit. Add a Single-mode APPROX_DISTINCT test, and let the clustered test helper take the aggregate function. Co-authored-by: Isaac --- .../src/aggregates/clustered_single_stream.rs | 24 ++++++++++++++++--- .../src/aggregates/clustered_test_util.rs | 16 +++++++++++-- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs index 4163925249b38..2e49b0e3b911a 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs @@ -762,9 +762,10 @@ mod tests { use super::*; use crate::aggregates::PhysicalGroupBy; use crate::aggregates::clustered_test_util::{ - assert_completed_groups_emitted_under_limit, assert_output_unchanged_under_limit, - assert_resources_exhausted_under_limit, distinct_key_batches, full_batches, - held_memory, single_key_batches, sum_over_sorted_input, + aggregate_over_sorted_input, assert_completed_groups_emitted_under_limit, + assert_output_unchanged_under_limit, assert_resources_exhausted_under_limit, + distinct_key_batches, full_batches, held_memory, single_key_batches, + sum_over_sorted_input, }; use crate::collect; use crate::test::TestMemoryExec; @@ -776,6 +777,7 @@ mod tests { GreedyMemoryPool, MemoryPool, PeakRecordingPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -874,6 +876,22 @@ mod tests { assert_output_unchanged_under_limit(&plan, 100 * 1024).await } + /// Like [`completed_groups_release_batch_capacity`], but `APPROX_DISTINCT` + /// also keeps the hashes of the last batch (64 KB for 8192 rows), which the + /// table must release too. The limit fits the active group, but not those + /// hashes. + #[tokio::test] + async fn completed_groups_release_accumulator_scratch() -> Result<()> { + let plan = aggregate_over_sorted_input( + AggregateMode::Single, + &["k"], + distinct_key_batches(), + approx_distinct_udaf(), + )?; + assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); + assert_output_unchanged_under_limit(&plan, 64 * 1024).await + } + /// A single key never completes a group before the input ends. The limit /// fits only the initial table. #[tokio::test] diff --git a/datafusion/physical-plan/src/aggregates/clustered_test_util.rs b/datafusion/physical-plan/src/aggregates/clustered_test_util.rs index 79e566843fa8c..851cc18502b88 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_test_util.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_test_util.rs @@ -18,6 +18,7 @@ use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::AggregateUDF; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -63,6 +64,17 @@ pub(super) fn sum_over_sorted_input( mode: AggregateMode, keys: &[&str], batches: Vec>>, +) -> Result { + aggregate_over_sorted_input(mode, keys, batches, sum_udaf()) +} + +/// Like [`sum_over_sorted_input`], but aggregates `v` with `udaf`. For `Final` +/// mode, the state of `udaf` must be a single `Int64` column. +pub(super) fn aggregate_over_sorted_input( + mode: AggregateMode, + keys: &[&str], + batches: Vec>>, + udaf: Arc, ) -> Result { let fields = keys .iter() @@ -76,9 +88,9 @@ pub(super) fn sum_over_sorted_input( .collect::>()?, ); let aggregate_exprs = vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + AggregateExprBuilder::new(Arc::clone(&udaf), vec![col("v", &schema)?]) .schema(Arc::clone(&schema)) - .alias("sum") + .alias(udaf.name()) .build()?, )];