Fix GroupValues retained memory accounting

datafusion/core/tests/memory_limit/mod.rs

@@ -123,7 +123,7 @@ async fn group_by_hash() {
         .with_query("select count(*) from t GROUP BY service, host, pod, container")
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "for PartialHashAggregateStream[0]",
+            "for FinalHashAggregateStream[0]",
         ])
         .with_memory_limit(1_000)
         .run()
@@ -747,7 +747,7 @@ async fn oom_grouped_hash_aggregate() {
         .with_query("SELECT COUNT(*), SUM(request_bytes) FROM t GROUP BY host")
         .with_expected_errors(vec![
             "Failed to allocate additional",
-            "for PartialHashAggregateStream[0]",
+            "for FinalHashAggregateStream[0]",
         ])
         .with_memory_limit(1_000)
         .run()

datafusion/core/tests/sql/aggregates/nested_nullability.rs

@@ -31,26 +31,28 @@
 //!
 //! [`Schema::contains`]: arrow::datatypes::Schema::contains
 
-use std::sync::Arc;
+use std::{num::NonZeroUsize, sync::Arc};
 
 use arrow::array::{BooleanArray, RecordBatch, StructArray, UInt32Array};
 use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
 use datafusion::datasource::MemTable;
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::physical_expr::aggregate::AggregateExprBuilder;
-use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy,
 };
 use datafusion::physical_plan::collect;
 use datafusion::physical_plan::expressions::col;
+use datafusion::physical_plan::{ExecutionPlan, displayable};
 use datafusion::prelude::*;
-use datafusion_common::Result;
+use datafusion_common::{Result, ScalarValue};
 use datafusion_execution::TaskContext;
-use datafusion_execution::memory_pool::FairSpillPool;
+use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool};
 use datafusion_execution::runtime_env::RuntimeEnvBuilder;
 use datafusion_functions_aggregate::array_agg::array_agg_udaf;
 
+use crate::helper::plan_metrics::{plan_spill_count, plan_spilled_bytes};
+
 /// Returns the fields of the struct column `b`: a single `colA Boolean`.
 ///
 /// `col_a_nullable` controls whether `colA` is declared nullable — the only
@@ -89,13 +91,19 @@ struct AggregateBatchesTest {
     /// If set, the context uses a [`FairSpillPool`] of this size (and a small
     /// batch size) so the aggregation is forced to spill.
     memory_limit: Option<usize>,
+    /// If set, fixes aggregate parallelism for deterministic memory pressure.
+    target_partitions: Option<usize>,
+    /// If set, test native DISTINCT aggregation rather than its group-by rewrite.
+    disable_single_distinct_to_groupby: bool,
 }
 
 impl AggregateBatchesTest {
     fn new() -> Self {
         Self {
             num_rows: 100,
             memory_limit: None,
+            target_partitions: None,
+            disable_single_distinct_to_groupby: false,
         }
     }
 
@@ -109,6 +117,16 @@ impl AggregateBatchesTest {
         self
     }
 
+    fn with_target_partitions(mut self, target_partitions: usize) -> Self {
+        self.target_partitions = Some(target_partitions);
+        self
+    }
+
+    fn without_single_distinct_to_groupby(mut self) -> Self {
+        self.disable_single_distinct_to_groupby = true;
+        self
+    }
+
     /// Runs `sql` against the table described above and asserts the result
     /// has one output row per group (i.e. [`Self::num_rows`] rows in total).
     async fn run(self, sql: &str) -> Result<()> {
@@ -138,22 +156,55 @@ impl AggregateBatchesTest {
 
         let ctx = match self.memory_limit {
             Some(limit) => {
+                // Include live consumers and peaks in any memory-pool failure.
+                // The FairSpillPool limit alone does not identify which concurrent
+                // spillable reservations divided its per-consumer allocation.
+                let memory_pool = TrackConsumersPool::new(
+                    FairSpillPool::new(limit),
+                    NonZeroUsize::new(10).unwrap(),
+                );
                 let runtime = RuntimeEnvBuilder::new()
-                    .with_memory_pool(Arc::new(FairSpillPool::new(limit)))
+                    .with_memory_pool(Arc::new(memory_pool))
                     .build_arc()?;
-                SessionContext::new_with_config_rt(
-                    SessionConfig::new().with_batch_size(100),
-                    runtime,
-                )
+                let mut config = SessionConfig::new().with_batch_size(100).set(
+                    "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
+                    &ScalarValue::Float64(Some(1.0)),
+                );
+                if let Some(target_partitions) = self.target_partitions {
+                    config = config.with_target_partitions(target_partitions);
+                }
+                SessionContext::new_with_config_rt(config, runtime)
             }
             None => SessionContext::new(),
         };
         ctx.register_table("t", Arc::new(table))?;
+        if self.disable_single_distinct_to_groupby {
+            assert!(ctx.remove_optimizer_rule("single_distinct_aggregation_to_group_by"));
+        }
 
-        let result = ctx.sql(sql).await?.collect().await?;
+        let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+        if self.disable_single_distinct_to_groupby {
+            let plan = displayable(plan.as_ref()).indent(true).to_string();
+            assert_eq!(
+                plan.matches("AggregateExec").count(),
+                1,
+                "expected native DISTINCT aggregation:\n{plan}"
+            );
+        }
+        let result = collect(Arc::clone(&plan), ctx.task_ctx()).await?;
 
         let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum();
         assert_eq!(total_rows, self.num_rows as usize);
+        if self.memory_limit.is_some() {
+            assert!(
+                plan_spill_count(plan.as_ref()) > 0,
+                "expected aggregation to spill"
+            );
+            assert!(
+                plan_spilled_bytes(plan.as_ref()) > 0,
+                "expected aggregation to spill bytes"
+            );
+        }
         Ok(())
     }
 }
@@ -176,7 +227,7 @@ async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> {
 async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> {
     AggregateBatchesTest::new()
         .with_num_rows(10_000)
-        .with_memory_limit(4_000_000)
+        .with_memory_limit(1_000_000)
         .run("SELECT a, array_agg(b) FROM t GROUP BY a")
         .await
 }
@@ -185,7 +236,10 @@ async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> {
 async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> {
     AggregateBatchesTest::new()
         .with_num_rows(10_000)
-        .with_memory_limit(4_000_000)
+        // One partition keeps the native aggregate's memory pressure deterministic.
+        .with_target_partitions(1)
+        .without_single_distinct_to_groupby()
+        .with_memory_limit(1_000_000)
         .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a")
         .await
 }

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

@@ -366,9 +366,16 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         let batch = RecordBatch::try_new(state_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
 
-        // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
-        // key/index buffers too so the memory reservation can be released
-        // before the batch is sorted for spilling.
+        // State emission should reset accumulators, but spill recovery must
+        // release every emitted allocation even for an accumulator that retains
+        // capacity. Rebuild the accumulator set before returning the state batch.
+        state.accumulators = state
+            .accumulators
+            .iter()
+            .map(HashAggregateAccumulator::empty_like)
+            .collect::<Result<_>>()?;
+        // Explicitly shrink key/index buffers too so the memory reservation can
+        // be released before the batch is sorted for spilling.
         state.group_values.clear_shrink(0);
         state.batch_group_indices.clear();
         state.batch_group_indices.shrink_to_fit();

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

@@ -341,9 +341,17 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?;
         debug_assert!(batch.num_rows() > 0);
 
-        // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
-        // key/index buffers too so the memory reservation can be released
-        // before the batch is passed downstream or sorted for spilling.
+        // State emission should reset accumulators, but spill recovery must
+        // release every emitted allocation even for an accumulator that retains
+        // capacity. Rebuild the accumulator set before returning the state batch.
+        self.buffer.accumulators = self
+            .buffer
+            .accumulators
+            .iter()
+            .map(AggregateAccumulator::empty_like)
+            .collect::<Result<_>>()?;
+        // Explicitly shrink key/index buffers too so the memory reservation can
+        // be released before the batch is passed downstream or sorted for spilling.
         self.buffer.group_values.clear_shrink(0);
         self.buffer.group_indices.clear();
         self.buffer.group_indices.shrink_to_fit();

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

@@ -21,6 +21,7 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, assert_eq_or_internal_err};
+use datafusion_expr::EmitTo;
 
 use crate::aggregates::group_values::{AccumulatorPhase, new_group_values};
 use crate::aggregates::order::GroupOrdering;
@@ -105,6 +106,65 @@ impl AggregateHashTable<PartialMarker> {
         })
     }
 
+    /// Starts a bounded-memory drain of partial aggregate states.
+    pub(in crate::aggregates) fn start_early_emit(&mut self) {
+        self.start_outputting();
+    }
+
+    /// Emits at most one output batch while releasing its groups from the table.
+    ///
+    /// Unlike terminal output, this must not materialize all states: early
+    /// emission can be triggered precisely because the complete state does not
+    /// fit in the memory pool. Once drained, rebuild an empty table so raw input
+    /// aggregation can resume.
+    pub(in crate::aggregates) fn next_early_emit_batch(
+        &mut self,
+    ) -> Result<Option<RecordBatch>> {
+        let state_schema = Arc::clone(&self.state_schema);
+        let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics);
+        let group_by_metrics = self.group_by_metrics.clone();
+        let AggregateHashTableState::Outputting(mut state) =
+            std::mem::replace(&mut self.state, AggregateHashTableState::Done)
+        else {
+            return Ok(None);
+        };
+
+        let emit_to = EmitTo::First(self.batch_size.min(state.group_values.len()));
+        let columns = group_by_metrics.time_emitting(|| {
+            let mut columns = state.group_values.emit(emit_to)?;
+            for (idx, acc) in state.accumulators.iter_mut().enumerate() {
+                columns.extend(accumulator_metrics.time(
+                    idx,
+                    AccumulatorPhase::State,
+                    || acc.state(emit_to),
+                )?);
+            }
+            Ok::<_, datafusion_common::DataFusionError>(columns)
+        })?;
+        let batch = RecordBatch::try_new(state_schema, columns)?;
+        debug_assert!(batch.num_rows() > 0);
+
+        if state.group_values.is_empty() {
+            let group_schema = state.group_by.group_schema(&self.input_schema)?;
+            let group_values = new_group_values(group_schema, &GroupOrdering::None)?;
+            let accumulators = state
+                .accumulators
+                .iter()
+                .map(HashAggregateAccumulator::empty_like)
+                .collect::<Result<Vec<_>>>()?;
+            self.state = AggregateHashTableState::Building(AggregateHashTableBuffer {
+                group_by: state.group_by,
+                group_values,
+                batch_group_indices: Vec::new(),
+                accumulators,
+            });
+        } else {
+            self.state = AggregateHashTableState::Outputting(state);
+        }
+
+        Ok(Some(batch))
+    }
+
     /// Partial aggregation consumes raw input rows and updates the table's
     /// partial-state accumulators.
     pub(in crate::aggregates) fn aggregate_batch(

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

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::mem::size_of;
 use std::sync::Arc;
 
 use crate::aggregates::group_values::multi_group_by::Nulls;
@@ -164,7 +165,7 @@ impl<const NULLABLE: bool> GroupColumn for BooleanGroupValueBuilder<NULLABLE> {
     }
 
     fn size(&self) -> usize {
-        self.buffer.capacity() / 8 + self.nulls.allocated_size()
+        size_of::<Self>() + self.buffer.capacity() / 8 + self.nulls.allocated_size()
     }
 
     fn build(self: Box<Self>) -> ArrayRef {

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

@@ -343,7 +343,8 @@ where
     }
 
     fn size(&self) -> usize {
-        self.buffer.capacity() * size_of::<u8>()
+        size_of::<Self>()
+            + self.buffer.capacity() * size_of::<u8>()
             + self.offsets.allocated_size()
             + self.nulls.allocated_size()
     }

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

@@ -28,6 +28,7 @@ use datafusion_common::utils::proxy::VecAllocExt;
 use datafusion_common::utils::split_vec_min_alloc;
 use datafusion_common::{Result, exec_datafusion_err};
 use datafusion_expr::GroupSelection;
+use std::mem::size_of;
 use std::sync::Arc;
 
 /// An implementation of [`GroupColumn`] for `FixedSizeBinary` values
@@ -212,7 +213,7 @@ impl GroupColumn for FixedSizeBinaryGroupValueBuilder {
     }
 
     fn size(&self) -> usize {
-        self.buffer.allocated_size() + self.nulls.allocated_size()
+        size_of::<Self>() + self.buffer.allocated_size() + self.nulls.allocated_size()
     }
 
     fn build(self: Box<Self>) -> ArrayRef {

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

@@ -40,6 +40,7 @@ use datafusion_common::utils::split_vec_min_alloc;
 use datafusion_common::{Result, internal_datafusion_err};
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use datafusion_expr::GroupSelection;
+use std::mem::size_of;
 use std::sync::Arc;
 
 /// A [`GroupColumn`] for `List<T>` (`O = i32`) and `LargeList<T>` (`O = i64`).
@@ -184,7 +185,8 @@ impl<O: OffsetSizeTrait> GroupColumn for ListGroupValueBuilder<O> {
     }
 
     fn size(&self) -> usize {
-        self.offsets.allocated_size()
+        size_of::<Self>()
+            + self.offsets.allocated_size()
             + self.outer_nulls.allocated_size()
             + self.child.size()
     }