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/dictionary.rs

@@ -30,7 +30,7 @@ use datafusion_expr::GroupSelection;
 use hashbrown::{HashMap, hash_table::HashTable};
 use std::marker::PhantomData;
 use std::mem::size_of;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use crate::aggregates::AGGREGATION_HASH_SEED;
 
@@ -58,9 +58,10 @@ pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync>
     val_to_inner: Vec<usize>,
     /// Reusable hash buffer for the dictionary values array.
     val_hashes: Vec<u64>,
-    /// The last `dict.values()` Arc hashed in `append_val`. When the incoming
-    /// values array is `ptr_eq` to this, `val_hashes` can be reused directly.
-    cached_values: Option<ArrayRef>,
+    /// Weak reference to the last `dict.values()` hashed in `append_val`.
+    /// This enables hash reuse when the array remains live without retaining an
+    /// input dictionary after its batch is released.
+    cached_values: Option<Weak<dyn Array>>,
     _phantom: PhantomData<K>,
 }
 
@@ -302,13 +303,13 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
             }
             Some(val_idx) => {
                 let dict_values = dict.values();
-                // check if the dictionary values array we are hashing was already seen.
-                // if its arc was already stored we dont need to rehash the entire array again
-                // if its new hash the entire array and store an arc ptr for future use
-                let cache_hit = self
-                    .cached_values
-                    .as_ref()
-                    .is_some_and(|c| Arc::ptr_eq(c, dict_values));
+                // Reuse hashes when the dictionary values array is still live
+                // and identical to the prior input; otherwise refresh the cache.
+                let cache_hit = self.cached_values.as_ref().is_some_and(|cached| {
+                    cached
+                        .upgrade()
+                        .is_some_and(|values| Arc::ptr_eq(&values, dict_values))
+                });
                 if !cache_hit {
                     self.val_hashes.clear();
                     self.val_hashes.resize(dict_values.len(), 0);
@@ -318,7 +319,7 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
                         &mut self.val_hashes,
                     )
                     .unwrap();
-                    self.cached_values = Some(Arc::clone(dict_values));
+                    self.cached_values = Some(Arc::downgrade(dict_values));
                 }
                 self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx])?
             }
@@ -670,6 +671,29 @@ mod tests {
         (0..buf.len()).map(|i| buf.get_bit(i)).collect()
     }
 
+    #[test]
+    fn hash_cache_does_not_retain_dictionary_values() {
+        let mut column = utf8_col();
+        let input = i32_dict(&[Some(0)], &[Some("retained")]);
+        let values = Arc::clone(input.as_dictionary::<Int32Type>().values());
+        let weak_values = Arc::downgrade(&values);
+        column.append_val(&input, 0).unwrap();
+        assert!(
+            column
+                .cached_values
+                .as_ref()
+                .is_some_and(|cached| cached.ptr_eq(&weak_values))
+        );
+        let size_with_input = column.size();
+
+        // The cache is only an identity hint: it must not keep an input
+        // dictionary allocation alive after its batch is released.
+        drop(values);
+        drop(input);
+        assert!(weak_values.upgrade().is_none());
+        assert_eq!(column.size(), size_with_input);
+    }
+
     fn all_true(len: usize) -> BooleanBufferBuilder {
         let mut buf = BooleanBufferBuilder::new(len);
         buf.append_n(len, true);