diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 01e9f3eaa71ee..1bbd3e22975eb 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -28,11 +28,10 @@ use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::normalize_float_zero; -use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupSelection}; use hashbrown::hash_table::HashTable; use log::debug; -use std::mem::size_of; use std::sync::Arc; /// A [`GroupValues`] making use of [`Rows`] @@ -60,9 +59,6 @@ pub struct GroupValuesRows { /// values: (hash, group_index) map: HashTable<(u64, usize)>, - /// The size of `map` in bytes - map_size: usize, - /// The actual group by values, stored in arrow [`Row`] format. /// `group_values[i]` holds the group value for group_index `i`. /// @@ -107,7 +103,6 @@ impl GroupValuesRows { schema, row_converter, map, - map_size: 0, group_values: None, hashes_buffer: Default::default(), rows_buffer, @@ -168,10 +163,10 @@ impl GroupValues for GroupValuesRows { group_values.push(group_rows.row(row)); // for hasher function, use precomputed hash value - self.map.insert_accounted( + self.map.insert_unique( + target_hash, (target_hash, group_idx), |(hash, _group_index)| *hash, - &mut self.map_size, ); group_idx } @@ -186,9 +181,11 @@ impl GroupValues for GroupValuesRows { fn size(&self) -> usize { let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0); + // `HashTable::allocation_size` reports the complete hashbrown allocation, + // which includes the entry array, control bytes, and trailing group layout. self.row_converter.size() + group_values_size - + self.map_size + + self.map.allocation_size() + self.rows_buffer.size() + self.hashes_buffer.allocated_size() } @@ -289,7 +286,6 @@ impl GroupValues for GroupValuesRows { }); 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); } @@ -443,8 +439,9 @@ pub(crate) fn encode_array_if_necessary( #[cfg(test)] mod tests { use super::*; - use arrow::array::{AsArray, ListArray}; + use arrow::array::{AsArray, Int32Array, ListArray, StringArray}; use arrow::datatypes::{Field, Int32Type, Schema}; + use std::mem::size_of; #[test] fn preserving_nested_row_values() -> Result<()> { @@ -484,4 +481,176 @@ mod tests { assert_eq!(groups, vec![3]); Ok(()) } + + #[test] + fn test_exact_hash_table_allocation_accounting_multi_column() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ])); + let mut group_values = GroupValuesRows::try_new(schema)?; + + // 1. Initial state: map is unallocated + assert_eq!(group_values.map.capacity(), 0); + assert_eq!(group_values.map.allocation_size(), 0); + assert_eq!( + group_values.size(), + group_values.row_converter.size() + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size() + ); + + // 2. Insert multi-column keys and observe table allocation growth + let num_distinct = 500; + let col_a = Arc::new(Int32Array::from_iter_values(0..num_distinct)) as ArrayRef; + let col_b = Arc::new(StringArray::from_iter_values( + (0..num_distinct).map(|i| format!("key_string_{i}")), + )) as ArrayRef; + + let mut groups = vec![]; + group_values.intern(&[col_a, col_b], &mut groups)?; + 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: capacity * size_of::<(u64, usize)>(). + let entry_only_bytes = group_values.map.capacity() * size_of::<(u64, usize)>(); + assert!( + table_bytes > entry_only_bytes, + "allocation_size ({table_bytes}) must exceed entry-only capacity ({entry_only_bytes})" + ); + + // GroupValuesRows::size() should accurately account for the map's allocation_size + let expected_size = group_values.row_converter.size() + + group_values + .group_values + .as_ref() + .map(|v| v.size()) + .unwrap_or(0) + + table_bytes + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size(); + assert_eq!(group_values.size(), expected_size); + + // 3. Partial emit (EmitTo::First) - table retains entries and capacity + let emit_count = 200; + let emitted = group_values.emit(EmitTo::First(emit_count))?; + 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(), + group_values.row_converter.size() + + group_values + .group_values + .as_ref() + .map(|v| v.size()) + .unwrap_or(0) + + table_bytes + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size() + ); + + // 4. Emit all (EmitTo::All) - map is cleared but capacity remains allocated for reuse + let remaining_count = group_values.len(); + let emitted_all = group_values.emit(EmitTo::All)?; + 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, + "clear_shrink should shrink table from {table_bytes} to {shrunken_table_bytes}" + ); + assert_eq!( + group_values.size(), + group_values.row_converter.size() + + group_values + .group_values + .as_ref() + .map(|v| v.size()) + .unwrap_or(0) + + shrunken_table_bytes + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size() + ); + + // 6. Reinsert new distinct values into the shrunken table + let new_col_a = + Arc::new(Int32Array::from_iter_values(num_distinct..num_distinct * 2)) + as ArrayRef; + let new_col_b = Arc::new(StringArray::from_iter_values( + (num_distinct..num_distinct * 2).map(|i| format!("key_string_{i}")), + )) as ArrayRef; + group_values.intern(&[new_col_a, new_col_b], &mut groups)?; + assert_eq!(groups.len(), num_distinct as usize); + assert!(group_values.map.allocation_size() > shrunken_table_bytes); + assert_eq!( + group_values.size(), + group_values.row_converter.size() + + group_values + .group_values + .as_ref() + .map(|v| v.size()) + .unwrap_or(0) + + group_values.map.allocation_size() + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size() + ); + + Ok(()) + } + + #[test] + fn test_exact_hash_table_allocation_accounting_nested_keys() -> Result<()> { + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let schema = Arc::new(Schema::new(vec![Field::new( + "nested", + DataType::List(field), + true, + )])); + let mut group_values = GroupValuesRows::try_new(schema)?; + + assert_eq!(group_values.map.capacity(), 0); + assert_eq!(group_values.map.allocation_size(), 0); + + let input = Arc::new(ListArray::from_iter_primitive::( + (0..256).map(|i| Some(vec![Some(i), Some(i * 2)])), + )) as ArrayRef; + let mut groups = vec![]; + group_values.intern(&[input], &mut groups)?; + assert_eq!(groups.len(), 256); + + let table_bytes = group_values.map.allocation_size(); + assert!(table_bytes > 0); + assert_eq!( + group_values.size(), + group_values.row_converter.size() + + group_values + .group_values + .as_ref() + .map(|v| v.size()) + .unwrap_or(0) + + table_bytes + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size() + ); + + // Shrink to 0: table should completely deallocate + group_values.clear_shrink(0); + assert_eq!(group_values.map.capacity(), 0); + assert_eq!(group_values.map.allocation_size(), 0); + + Ok(()) + } }