fix: release per-batch accumulator scratch after a clustered partial emit under memory pressure

datafusion/expr-common/src/groups_accumulator.rs

@@ -381,6 +381,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)]

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs

@@ -666,6 +666,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],

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -642,6 +642,11 @@ impl GroupsAccumulator for HllGroupsAccumulator {
             + self.allocated_bytes
             + self.hashes.capacity() * size_of::<u64>()
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.groups.shrink_to_fit();
+        self.hashes = Vec::new();
+    }
 }
 
 impl Debug for ApproxDistinct {

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -829,6 +829,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,

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs

@@ -299,6 +299,19 @@ impl<AggrMode> ClusteredAggregateTable<AggrMode> {
             + self.buffer.group_indices.allocated_size()
     }
 
+    /// Releases the capacity that the buffer holds beyond its stored groups,
+    /// 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 {
         ClusteredAggregateTableMetrics {
             group_by: self.group_by_metrics.clone(),

datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs

@@ -87,7 +87,9 @@ use crate::{SendableRecordBatchStream, metrics};
 ///
 /// If the reservation after an input batch fails, `ClusteredPartialAggregateStream`
 /// still emits the completed groups, and returns the error only when no group is
-/// complete.
+/// complete. After the emit it releases the capacity the table grew for the batch,
+/// such as hash table slots and per-batch scratch buffers, so only the active
+/// group stays reserved.
 ///
 /// ## Partial group clustering
 ///
@@ -262,6 +264,8 @@ impl ClusteredPartialAggregateStream {
     /// 3. Prepare output:
     ///    - A prefix of groups is complete: materialize the entire prefix once,
     ///      retaining the input and active groups to resume aggregation.
+    ///    - On memory pressure with full group clustering, materialize the
+    ///      completed prefix and release the capacity the table grew for it.
     ///    - On memory pressure with partial group clustering, materialize all
     ///      current states, including incomplete groups, and reset the table.
     ///    - At EOF, materialize all remaining states and prepare to output.
@@ -319,11 +323,17 @@ 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.
-                    // The final stage can merge incomplete states emitted here.
                     if matches!(self.table.group_clustering(), GroupClustering::Full(_)) {
-                        Some(self.table.take_completed_state_batch()?.ok_or(oom)?)
+                        // Only the active group remains after the emit. Release the
+                        // capacity the table grew for this batch so that reserving
+                        // the table can succeed.
+                        let batch =
+                            self.table.take_completed_state_batch()?.ok_or(oom)?;
+                        self.table.shrink_to_fit();
+                        Some(batch)
                     } else {
+                        // Partial group clustering may have an unbounded active key range.
+                        // The final stage can merge incomplete states emitted here.
                         Some(self.table.take_state_batch()?.ok_or(oom)?)
                     }
                 }
@@ -421,79 +431,211 @@ impl Outputting {
 
 #[cfg(test)]
 mod tests {
+    use std::ops::Range;
+
     use super::*;
     use crate::aggregates::PhysicalGroupBy;
     use crate::test::TestMemoryExec;
-    use arrow::array::Int64Array;
-    use arrow::datatypes::{DataType, Field, Schema};
+    use arrow::array::{ArrayRef, Float64Array, Int64Array, ListArray};
+    use arrow::datatypes::{DataType, Field, Int64Type, 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_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;
 
-    /// 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<dyn MemoryPool>) -> Result<(Vec<RecordBatch>, 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)])
+    /// One batch per range. Every key column and the value column `v` hold the
+    /// range values, as the given key types.
+    fn batches(
+        key_types: &[DataType],
+        ranges: &[Range<i64>],
+    ) -> Result<Vec<RecordBatch>> {
+        let mut fields: Vec<_> = key_types
+            .iter()
+            .enumerate()
+            .map(|(i, data_type)| Field::new(format!("k{i}"), data_type.clone(), false))
+            .collect();
+        fields.push(Field::new("v", DataType::Int64, false));
+        let schema = Arc::new(Schema::new(fields));
+        ranges
+            .iter()
+            .map(|range| {
+                let int64: ArrayRef =
+                    Arc::new(Int64Array::from_iter_values(range.clone()));
+                let mut columns = key_types
+                    .iter()
+                    .map(|data_type| match data_type {
+                        DataType::Int64 => Arc::clone(&int64),
+                        DataType::Float64 => Arc::new(Float64Array::from_iter_values(
+                            range.clone().map(|v| v as f64),
+                        )),
+                        DataType::List(_) => {
+                            Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(
+                                range.clone().map(|v| Some([Some(v)])),
+                            ))
+                        }
+                        _ => unreachable!("unsupported key type {data_type}"),
+                    })
+                    .collect::<Vec<_>>();
+                columns.push(int64);
+                Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?)
             })
-            .collect::<Result<Vec<_>, _>>()?;
+            .collect()
+    }
+
+    /// Runs a partial `SUM(v) GROUP BY k0, ..` over `batches`, which are sorted
+    /// on every key.
+    async fn run(
+        batches: Vec<RecordBatch>,
+        pool: Arc<dyn MemoryPool>,
+    ) -> Result<Vec<RecordBatch>> {
+        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<RecordBatch>,
+        aggregates: Vec<Arc<AggregateFunctionExpr>>,
+        pool: Arc<dyn MemoryPool>,
+    ) -> Result<Vec<RecordBatch>> {
+        let schema = batches[0].schema();
+        let keys = (0..schema.fields().len() - 1)
+            .map(|i| Ok((col(&format!("k{i}"), &schema)?, format!("k{i}"))))
+            .collect::<Result<Vec<_>>>()?;
+        let ordering = LexOrdering::new(
+            keys.iter()
+                .map(|(expr, _)| PhysicalSortExpr::new_default(Arc::clone(expr))),
+        )
+        .unwrap();
         let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)?
-            .try_with_sort_information(vec![
-                LexOrdering::new([PhysicalSortExpr::new_default(col("k", &schema)?)])
-                    .unwrap(),
-            ])?;
+            .try_with_sort_information(vec![ordering])?;
         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("k", &schema)?])
-                    .schema(Arc::clone(&schema))
-                    .alias("sum")
-                    .build()?,
-            )],
-            vec![None],
+            PhysicalGroupBy::new_single(keys),
+            aggregates.clone(),
+            vec![None; aggregates.len()],
             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))
+                    .with_memory_pool(pool)
                     .build_arc()?,
             ),
         );
-        let mut stream =
+        let 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))
+        crate::common::collect(stream).await
     }
 
-    /// 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.
+    /// Every row is a distinct group, and each batch shares its first key with
+    /// the previous one, so the active group spans the batches. Aggregating a
+    /// batch grows the table for about 8192 groups, for example the single-key
+    /// hash table (~272 KB) and the `group_indices` scratch (64 KB). After the
+    /// completed groups are emitted, the table must release that capacity, or
+    /// reserving it fails the limit although the active group needs a few KB.
+    /// The output 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));
+        let list = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true)));
+        for key_types in [
+            // `GroupValuesPrimitive`
+            vec![DataType::Int64],
+            // `GroupValuesClustered`
+            vec![DataType::Int64, DataType::Int64],
+            // `GroupValuesColumn`
+            vec![DataType::Float64, DataType::Float64],
+            // `GroupValuesRows`
+            vec![DataType::Int64, list],
+        ] {
+            let input = || batches(&key_types, &[0..8193, 8192..16385, 16384..24576]);
+            let expected =
+                run(input()?, Arc::new(UnboundedMemoryPool::default())).await?;
+            let actual = run(input()?, Arc::new(GreedyMemoryPool::new(100 * 1024)))
+                .await
+                .map_err(|e| e.context(format!("key types {key_types:?}")))?;
+            assert_eq!(batches_to_string(&actual), batches_to_string(&expected));
+        }
+        Ok(())
+    }
+
+    /// 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<Vec<Arc<AggregateFunctionExpr>>>;
+        fn build(
+            udaf: Arc<AggregateUDF>,
+            args: &[&str],
+            schema: &SchemaRef,
+        ) -> Result<Arc<AggregateFunctionExpr>> {
+            let alias = format!("{}({})", udaf.name(), args.join(", "));
+            let args = args
+                .iter()
+                .map(|name| col(name, schema))
+                .collect::<Result<Vec<_>>>()?;
+            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(())
     }
 }

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -133,6 +133,16 @@ pub trait GroupValues: Send {
 
     /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage)
     fn clear_shrink(&mut self, num_rows: usize);
+
+    /// Keeps the stored groups but releases capacity they do not need, such as
+    /// hash table slots and per-batch scratch buffers, so that [`Self::size`]
+    /// reflects the stored groups. For example, after [`EmitTo::First`] removes
+    /// most groups, the hash table may keep the slots it grew for them.
+    ///
+    /// The default does nothing.
+    ///
+    /// [`EmitTo::First`]: datafusion_expr::EmitTo::First
+    fn shrink_to_fit(&mut self) {}
 }
 
 /// Return a specialized implementation of [`GroupValues`] for the given schema.

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs

@@ -181,6 +181,12 @@ impl GroupValues for GroupValuesClustered {
         true
     }
 
+    fn shrink_to_fit(&mut self) {
+        // Scratch for one `intern` call.
+        self.new_groups.clear();
+        self.new_groups.shrink_to_fit();
+    }
+
     fn clear_shrink(&mut self, num_rows: usize) {
         self.columns = GroupValuesColumn::<true>::build_group_columns(&self.schema)
             .expect("schema validated by GroupValuesClustered::try_new");