diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index eaeb83e0cd712..816b020c39fb9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/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.map.capacity() * size_of::<(usize, u64)>() - + self.values.allocated_size() + size_of::() + 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::::new(DataType::Int32); + let self_size = size_of::>(); + + // 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::::new(DataType::Int32); diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 3c57b7fade5d3..79d1213ce4316 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/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 ( ) ---- -04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spill_count=7,] +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- 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 ( ) ---- -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=[spill_count=7,] +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=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- 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 ( ) ---- -04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[spill_count=7,] +04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- 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 ( ) ---- -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=[spill_count=7,] +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=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- 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 ( ) ---- -04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[avg(t.v)], metrics=[spill_count=7,] +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[avg(t.v)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- 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 ( ) ---- -04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[array_agg(t.v)], metrics=[spill_count=7,] +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[array_agg(t.v)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- Case G: partial/final aggregation under memory limit --- @@ -232,7 +234,7 @@ FROM ( ) ---- -06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spilled_rows=K,] +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spill_count=, spilled_bytes=, spilled_rows= K,] 08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[early_emit_count=] diff --git a/datafusion/sqllogictest/test_files/explain_analyze.slt b/datafusion/sqllogictest/test_files/explain_analyze.slt index 6fb01e1e270e8..f3ba4d6afc6a2 100644 --- a/datafusion/sqllogictest/test_files/explain_analyze.slt +++ b/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=[]