diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 2836b3177eb05..e99dda9b6a791 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -49,7 +49,7 @@ use datafusion_common::float_semantics::{has_float_leaf, normalize_floats}; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::{Result, not_impl_err}; -use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; @@ -194,9 +194,6 @@ pub struct GroupValuesColumn { /// map: HashTable<(u64, GroupIndexView)>, - /// The size of `map` in bytes - map_size: usize, - /// The lists for group indices with the same hash value /// /// It is possible that hash value collision exists, @@ -307,7 +304,6 @@ impl GroupValuesColumn { group_index_lists: Vec::new(), emit_group_index_list_buffer: Vec::new(), vectorized_operation_buffers: VectorizedOperationBuffers::default(), - map_size: 0, group_values, hashes_buffer: Default::default(), random_state: crate::aggregates::AGGREGATION_HASH_SEED, @@ -448,10 +444,10 @@ impl GroupValuesColumn { } // for hasher function, use precomputed hash value - self.map.insert_accounted( + self.map.insert_unique( + target_hash, (target_hash, GroupIndexView::new_inlined(group_idx as u64)), |(hash, _group_index)| *hash, - &mut self.map_size, ); group_idx } @@ -568,10 +564,10 @@ impl GroupValuesColumn { // Insert the `group index view` and its hash into `map` // for hasher function, use precomputed hash value - self.map.insert_accounted( + self.map.insert_unique( + target_hash, (target_hash, group_index_view), |(hash, _)| *hash, - &mut self.map_size, ); // Add row index to `vectorized_append_row_indices` @@ -1096,9 +1092,12 @@ impl GroupValues for GroupValuesColumn { .iter() .map(VecAllocExt::allocated_size) .sum::(); + // `HashTable::allocation_size` reports the complete retained allocation — + // including hashbrown control bytes and trailing layout — which is exactly + // what we want here. This follows the same approach as `ArrowBytesMap::size()`. size_of::() + group_values_size - + self.map_size + + self.map.allocation_size() + group_index_lists_size + self.emit_group_index_list_buffer.allocated_size() + self.vectorized_operation_buffers.size() @@ -1244,7 +1243,6 @@ impl GroupValues for GroupValuesColumn { .expect("schema previously validated in try_new"); self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared - self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); self.hashes_buffer.clear(); self.hashes_buffer.shrink_to(num_rows); @@ -1294,8 +1292,9 @@ mod tests { compute::{concat_batches, take}, util::pretty::pretty_format_batches, }; - use datafusion_common::{Result, utils::proxy::HashTableAllocExt}; + use datafusion_common::Result; use datafusion_expr::{EmitTo, GroupSelection}; + use std::mem::size_of; use crate::aggregates::group_values::{ GroupValues, multi_group_by::GroupValuesColumn, @@ -2764,10 +2763,10 @@ mod tests { group_index: u64, ) { let group_index_view = GroupIndexView::new_inlined(group_index); - group_values.map.insert_accounted( + group_values.map.insert_unique( + hash_key, (hash_key, group_index_view), |(hash, _)| *hash, - &mut group_values.map_size, ); } @@ -2779,10 +2778,159 @@ mod tests { let list_offset = group_values.group_index_lists.len(); let group_index_view = GroupIndexView::new_non_inlined(list_offset as u64); group_values.group_index_lists.push(group_indices); - group_values.map.insert_accounted( + group_values.map.insert_unique( + hash_key, (hash_key, group_index_view), |(hash, _)| *hash, - &mut group_values.map_size, + ); + } + + // ----------------------------------------------------------------------- + // Tests for exact hash-table allocation accounting (issue #25736) + // ----------------------------------------------------------------------- + + /// After enough groups are inserted to force a hashbrown table growth, + /// `size()` must include the full `allocation_size()` of the map — control + /// bytes and trailing layout included — not just an entry-capacity estimate. + #[test] + fn map_allocation_size_included_in_size_after_growth() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + // Insert enough distinct rows to trigger at least one table resize. + let n: usize = 128; + let keys: Vec = (0..n as i64).collect(); + let col = Arc::new(Int64Array::from(keys)) as ArrayRef; + let mut groups = vec![]; + gv.intern(&[col], &mut groups).unwrap(); + + let reported = gv.size(); + let map_alloc = gv.map.allocation_size(); + + // The reported size must be at least as large as the raw map allocation. + assert!( + reported >= map_alloc, + "size() ({reported}) must be >= map.allocation_size() ({map_alloc})" + ); + + // And the map must actually hold a non-trivial allocation after growth. + assert!( + map_alloc > 0, + "map.allocation_size() should be > 0 after inserting {n} groups" + ); + } + + /// When hash collisions are forced, the collision-chain list allocation + /// (`group_index_lists`) must be accounted for independently of the + /// hash-table term. The two must not be conflated. + #[test] + fn collision_chain_size_is_separate_from_map_allocation() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + // Manually insert two entries with the same hash key to force a + // collision chain (non-inlined GroupIndexView). + insert_inline_group_index_view(&mut gv, 42, 0); + insert_non_inline_group_index_view(&mut gv, 42, vec![0, 1]); + + let map_alloc = gv.map.allocation_size(); + let chain_size: usize = gv + .group_index_lists + .iter() + .map(|list| list.capacity() * size_of::()) + .sum(); + + // Both terms must be individually non-zero. + assert!(map_alloc > 0, "map allocation must be > 0"); + assert!(chain_size > 0, "collision-chain allocation must be > 0"); + + // The reported total must include both independently. + let reported = gv.size(); + assert!( + reported >= map_alloc, + "size() must cover the map allocation" + ); + } + + /// After a partial emit (`EmitTo::First`), the map retains its allocated + /// capacity even though logical entries were removed. `size()` must still + /// reflect the retained allocation, not drop to zero. + #[test] + fn size_reflects_retained_capacity_after_partial_emit() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + let keys: Vec = (0..64i64).collect(); + let col = Arc::new(Int64Array::from(keys)) as ArrayRef; + let mut groups = vec![]; + gv.intern(&[col], &mut groups).unwrap(); + + let size_before = gv.size(); + let map_alloc_before = gv.map.allocation_size(); + + // Emit only the first 32 groups; the table allocation should be retained. + gv.emit(EmitTo::First(32)).unwrap(); + + let size_after = gv.size(); + let map_alloc_after = gv.map.allocation_size(); + + // Capacity is not released by a partial emit — allocation should persist. + assert!( + map_alloc_after > 0, + "map allocation should be retained after partial emit, got {map_alloc_after}" + ); + assert!( + size_after > 0, + "size() should remain > 0 after partial emit, got {size_after}" + ); + + // The allocation reported before and after should both be real and + // consistent with what the map actually holds. + assert_eq!( + map_alloc_before, map_alloc_after, + "partial emit must not shrink the map allocation (before={map_alloc_before}, after={map_alloc_after})" + ); + let _ = size_before; // used implicitly via the assertions above + } + + /// After a full emit followed by re-use, `size()` must still produce + /// correct group values and the map allocation is reported accurately. + #[test] + fn reuse_after_full_emit_produces_correct_groups() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + // First pass: intern 4 distinct keys. + let first_keys: Vec = vec![10, 20, 30, 40]; + let col1 = Arc::new(Int64Array::from(first_keys.clone())) as ArrayRef; + let mut groups = vec![]; + gv.intern(&[col1], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 3]); + + // Full emit clears the map. + let emitted = gv.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 1); + + // Second pass after reuse: intern 3 different keys, numbering restarts from 0. + let second_keys: Vec = vec![100, 200, 300]; + let col2 = Arc::new(Int64Array::from(second_keys)) as ArrayRef; + gv.intern(&[col2], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2]); + + // size() must be consistent with the live map allocation. + let reported = gv.size(); + let map_alloc = gv.map.allocation_size(); + assert!( + reported >= map_alloc, + "size() ({reported}) must cover map.allocation_size() ({map_alloc}) after reuse" ); } }