fix(physical-plan): replace approximate map_size accounting with exac…

datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs

@@ -31,7 +31,6 @@ use half::f16;
 use hashbrown::hash_table::HashTable;
 #[cfg(not(feature = "force_hash_collisions"))]
 use std::hash::BuildHasher;
-use std::mem::size_of;
 use std::sync::Arc;
 
 /// A trait to allow hashing of floating point numbers
@@ -178,9 +177,7 @@ where
     }
 
     fn size(&self) -> usize {
-        size_of::<Self>()
-            + self.map.capacity() * size_of::<(usize, u64)>()
-            + self.values.allocated_size()
+        size_of::<Self>() + self.map.allocation_size() + self.values.allocated_size()
     }
 
     fn is_empty(&self) -> bool {
@@ -288,6 +285,7 @@ mod tests {
     use arrow::array::{ArrayRef, Int32Array};
     use arrow::datatypes::DataType;
     use datafusion_expr::EmitTo;
+    use std::mem::size_of;
     use std::sync::Arc;
 
     /// Mirror of the `EmitTo::take_needed` regression test, applied to the
@@ -328,6 +326,89 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_exact_hash_table_allocation_accounting() {
+        let mut group_values = GroupValuesPrimitive::<Int32Type>::new(DataType::Int32);
+        let self_size = size_of::<GroupValuesPrimitive<Int32Type>>();
+
+        // 1. Initial state: the constructor calls HashTable::with_capacity(128),
+        //    so the map already has a nonzero retained allocation. Capture it
+        //    and verify that size() accounts for it correctly.
+        let initial_map_bytes = group_values.map.allocation_size();
+        assert!(group_values.map.capacity() >= 128);
+        assert!(initial_map_bytes > 0);
+        assert_eq!(
+            group_values.size(),
+            self_size + group_values.values.allocated_size() + initial_map_bytes,
+        );
+
+        // 2. Insert keys and observe table allocation growth
+        let num_distinct = 500_i32;
+        let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..num_distinct));
+        let mut groups = vec![];
+        group_values.intern(&[arr], &mut groups).unwrap();
+        assert_eq!(groups.len(), num_distinct as usize);
+
+        let table_bytes = group_values.map.allocation_size();
+        assert!(group_values.map.capacity() >= num_distinct as usize);
+        assert!(table_bytes > 0);
+
+        // Hashbrown's layout contains entry slots, control bytes, and trailing group bytes.
+        // It must strictly exceed the naive entry-only size
+        let entry_only_bytes = group_values.map.capacity() * size_of::<(usize, u64)>();
+        assert!(table_bytes > entry_only_bytes);
+
+        assert_eq!(
+            group_values.size(),
+            self_size + group_values.values.allocated_size() + table_bytes,
+        );
+
+        // 3. Partial emit (EmitTo::First)
+        let emit_count = 200;
+        let emitted = group_values.emit(EmitTo::First(emit_count)).unwrap();
+        assert_eq!(emitted[0].len(), emit_count);
+        assert_eq!(group_values.len(), (num_distinct as usize) - emit_count);
+
+        // Retained capacity is still charged until the table releases it
+        assert_eq!(group_values.map.allocation_size(), table_bytes);
+        assert_eq!(
+            group_values.size(),
+            self_size + group_values.values.allocated_size() + table_bytes,
+        );
+
+        // 4. Emit all (EmitTo::All)
+        let remaining_count = group_values.len();
+        let emitted_all = group_values.emit(EmitTo::All).unwrap();
+        assert_eq!(emitted_all[0].len(), remaining_count);
+        assert_eq!(group_values.len(), 0);
+
+        // Clear preserves table capacity
+        assert_eq!(group_values.map.allocation_size(), table_bytes);
+
+        // 5. clear_shrink: releases/shrinks map capacity
+        group_values.clear_shrink(16);
+        let shrunken_table_bytes = group_values.map.allocation_size();
+        assert!(shrunken_table_bytes < table_bytes);
+        assert_eq!(
+            group_values.size(),
+            self_size + group_values.values.allocated_size() + shrunken_table_bytes,
+        );
+
+        // 6. Reinsert new distinct values into the shrunken table
+        let new_arr: ArrayRef = Arc::new(Int32Array::from_iter_values(
+            num_distinct..(num_distinct * 2),
+        ));
+        group_values.intern(&[new_arr], &mut groups).unwrap();
+        assert_eq!(groups.len(), num_distinct as usize);
+        assert!(group_values.map.allocation_size() > shrunken_table_bytes);
+        assert_eq!(
+            group_values.size(),
+            self_size
+                + group_values.values.allocated_size()
+                + group_values.map.allocation_size(),
+        );
+    }
+
     #[test]
     fn clear_shrink_reset_null_group() -> Result<()> {
         let mut gv = GroupValuesPrimitive::<Int32Type>::new(DataType::Int32);

datafusion/sqllogictest/test_files/aggregate_memory_spill.slt

@@ -50,7 +50,9 @@ FROM (
 ----
 100000 5000050000
 
-# Assert spill happened, the `spill_count` metric must be > 0
+# Assert spill happened: the aggregate spills all 100k rows, so
+# spilled_rows is formatted as "100.0 K". This value cannot appear
+# when spill_count=0 because spilled_rows would then be 0 (no K suffix).
 query TT
 EXPLAIN ANALYZE
 SELECT count(*), sum(total)
@@ -61,7 +63,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=100.0 K,<slt:ignore>]
 <slt:ignore>
 
 # --- Case B: multi-column GROUP BY (is_single() = false) ---
@@ -76,7 +78,7 @@ FROM (
 ----
 100000 5000050000
 
-# Assert spill happened, the `spill_count` metric must be > 0
+# Assert spill happened (same reasoning as Case A).
 query TT
 EXPLAIN ANALYZE
 SELECT count(*), sum(total)
@@ -87,7 +89,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000), v@0 * 13 % 100000 as t.v * Int64(13) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000), v@0 * 13 % 100000 as t.v * Int64(13) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=100.0 K,<slt:ignore>]
 <slt:ignore>
 
 # --- Case C: DISTINCT aggregate under memory limit ---
@@ -102,7 +104,7 @@ FROM (
 ----
 100000 100000
 
-# Assert spill happened, the `spill_count` metric must be > 0
+# Assert spill happened (same reasoning as Case A).
 query TT
 EXPLAIN ANALYZE
 SELECT count(*), sum(d)
@@ -113,7 +115,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=100.0 K,<slt:ignore>]
 <slt:ignore>
 
 # --- Case D: multiple aggregates (sum/min/max) under memory limit ---
@@ -128,7 +130,7 @@ FROM (
 ----
 100000 5000050000 1 100000
 
-# Assert spill happened, the `spill_count` metric must be > 0
+# Assert spill happened (same reasoning as Case A).
 query TT
 EXPLAIN ANALYZE
 SELECT count(*), sum(s), min(mn), max(mx)
@@ -139,7 +141,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v), min(t.v), max(t.v)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v), min(t.v), max(t.v)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=100.0 K,<slt:ignore>]
 <slt:ignore>
 
 # --- Case E: avg() aggregate (Float64 output) under memory limit ---
@@ -154,7 +156,7 @@ FROM (
 ----
 100000 1 100000
 
-# Assert spill happened, the `spill_count` metric must be > 0
+# Assert spill happened (same reasoning as Case A).
 query TT
 EXPLAIN ANALYZE
 SELECT count(*), min(a), max(a)
@@ -165,7 +167,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[avg(t.v)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[avg(t.v)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=100.0 K,<slt:ignore>]
 <slt:ignore>
 
 # --- Case F: array_agg() aggregate (growable state) under memory limit ---
@@ -180,7 +182,7 @@ FROM (
 ----
 100000 100000
 
-# Assert spill happened, the `spill_count` metric must be > 0
+# Assert spill happened (same reasoning as Case A).
 query TT
 EXPLAIN ANALYZE
 SELECT count(*), sum(l)
@@ -191,7 +193,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[array_agg(t.v)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[array_agg(t.v)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=100.0 K,<slt:ignore>]
 <slt:ignore>
 
 # --- Case G: partial/final aggregation under memory limit ---
@@ -232,7 +234,7 @@ FROM (
 )
 ----
 <slt:ignore>
-06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spilled_rows=<slt:ignore>K,<slt:ignore>]
+06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=<slt:ignore>, spilled_bytes=<slt:ignore>, spilled_rows=<slt:ignore> K,<slt:ignore>]
 <slt:ignore>
 08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>early_emit_count=<slt:ignore>]
 <slt:ignore>

datafusion/sqllogictest/test_files/explain_analyze.slt

@@ -500,7 +500,7 @@ GROUP BY k;
 ----
 Plan with Metrics
 01)ProjectionExec: expr=[k@0 as k, count(Int64(1))@1 as count(*)], metrics=[output_bytes=1056.0 B]
-02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=9.4 KB]
+02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=10.9 KB]
 03)----ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B]
 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]