From a1a3d5fb007171fe55729478332675f412cd4316 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Wed, 28 Oct 2026 19:11:23 -0400 Subject: [PATCH 1/5] refactor: add reusable caching array hasher Reuse dictionary value hashes by column while the values Arc is unchanged, without retaining input dictionaries. Expose cache state for callers that account for memory and document the public hashing contract. --- datafusion/common/src/hash_utils.rs | 346 ++++++++++++++++++++++++++++ 1 file changed, 346 insertions(+) diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index c9cc48a11dd21..59079836dc071 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -17,6 +17,10 @@ use foldhash::fast::FixedState; #[cfg(not(feature = "force_hash_collisions"))] use itertools::Itertools; use std::hash::{BuildHasher, Hash, Hasher}; +use std::mem::size_of; +#[cfg(not(feature = "force_hash_collisions"))] +use std::sync::Arc; +use std::sync::Weak; /// [`RandomState`] is optimized for speed and suitable for hash tables and /// bloom filters. [`QualityRandomState`] is optimized for statistical quality @@ -28,6 +32,146 @@ use std::hash::{BuildHasher, Hash, Hasher}; pub type RandomState = FixedState; pub type QualityRandomState = foldhash::quality::FixedState; +type DictionaryHashEntry = (Weak, Vec); + +/// Hashes arrays, reusing dictionary value hashes by column when possible. +#[derive(Debug)] +pub struct CachingArrayHasher { + /// Stored with the cache so value hashes cannot cross hash seeds. + random_state: RandomState, + /// One entry per column. Weak identity avoids retaining dictionary values. + dictionaries: Vec>, +} + +impl CachingArrayHasher { + /// Creates one cache slot per column in `schema`. + /// + /// For each dictionary column, the cache tracks the values array with a + /// weak pointer, so it does not keep the values alive, and owns a reusable + /// value-hash buffer. [`Self::create_hashes`] reuses those hashes while the + /// values `Arc` is unchanged, recomputes them when it changes, and scatters + /// them through the dictionary keys into the row hashes. Non-dictionary + /// columns use the regular array hashing path. + pub fn new(schema: &Schema, random_state: RandomState) -> Self { + Self { + random_state, + dictionaries: schema.fields().iter().map(|_| None).collect(), + } + } + + /// Hashes `arrays`, reusing dictionary value hashes while each column's + /// values array is unchanged. + /// + /// `arrays` must contain exactly one array for each field in the schema + /// passed to [`Self::new`], in the same order. Every array must have + /// `hashes_buffer.len()` rows. + /// + /// # Errors + /// + /// Returns an error if the number of arrays does not match the schema or + /// if an array cannot be hashed. + pub fn create_hashes<'a, I, T>( + &mut self, + arrays: I, + hashes_buffer: &'a mut [u64], + ) -> Result<&'a mut [u64]> + where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, + T: AsDynArray, + { + let arrays = arrays.into_iter(); + if arrays.len() != self.dictionaries.len() { + return _internal_err!( + "CachingArrayHasher expected {} arrays, got {}", + self.dictionaries.len(), + arrays.len() + ); + } + for (index, array) in arrays.enumerate() { + self.hash_array(index, array.as_dyn_array(), hashes_buffer, index != 0)?; + } + Ok(hashes_buffer) + } + + fn hash_array( + &mut self, + index: usize, + array: &dyn Array, + hashes_buffer: &mut [u64], + rehash: bool, + ) -> Result<()> { + #[cfg(feature = "force_hash_collisions")] + { + let _ = index; + hash_single_array(array, &self.random_state, hashes_buffer, rehash)?; + } + + #[cfg(not(feature = "force_hash_collisions"))] + { + let dictionary = &mut self.dictionaries[index]; + if !matches!(array.data_type(), DataType::Dictionary(_, _)) { + *dictionary = None; + hash_single_array(array, &self.random_state, hashes_buffer, rehash)?; + return Ok(()); + } + + downcast_dictionary_array! { + array => { + let values = array.values(); + if !dictionary.as_ref().is_some_and(|(cached, _)| { + std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values)) + }) { + let mut hashes = dictionary + .take() + .map(|(_, hashes)| hashes) + .unwrap_or_default(); + hashes.resize(values.len(), 0); + // Fresh hashing leaves null value slots unchanged. + hashes.fill(0); + hash_single_array( + values.as_ref(), + &self.random_state, + &mut hashes, + false, + )?; + *dictionary = Some((Arc::downgrade(values), hashes)); + } + dispatch_dictionary_scatter( + array, + &dictionary.as_ref().unwrap().1, + hashes_buffer, + rehash, + ); + } + _ => unreachable!(), + }; + } + + Ok(()) + } + + /// Returns the current dictionary values and hashes for `index`. + /// + /// Returns `None` if `index` is out of range, its cache slot does not + /// contain a dictionary, or the dictionary values have been dropped. + pub fn dictionary_hashes(&self, index: usize) -> Option<(ArrayRef, &[u64])> { + let (values, hashes) = self.dictionaries.get(index)?.as_ref()?; + Some((values.upgrade()?, hashes)) + } + + /// Returns the allocated size of the cache. + pub fn size(&self) -> usize { + self.dictionaries.capacity() * size_of::>() + + self + .dictionaries + .iter() + .flatten() + .map(|entry| entry.1.capacity() * size_of::()) + .sum::() + } +} + /// Fixed quality hash state used by HyperLogLog sketches. /// /// The seed is part of the HLL wire/storage semantics: serialized sketches only @@ -1357,6 +1501,208 @@ mod tests { } } + #[cfg(not(feature = "force_hash_collisions"))] + fn string_dictionary(keys: Vec, values: &ArrayRef) -> ArrayRef { + Arc::new(DictionaryArray::::new( + Int8Array::from(keys), + Arc::clone(values), + )) + } + + #[cfg(not(feature = "force_hash_collisions"))] + fn caching_hasher( + columns: &[ArrayRef], + random_state: &RandomState, + ) -> CachingArrayHasher { + let fields = columns + .iter() + .enumerate() + .map(|(index, column)| { + Field::new(format!("column_{index}"), column.data_type().clone(), true) + }) + .collect::>(); + CachingArrayHasher::new(&Schema::new(fields), random_state.clone()) + } + + #[cfg(not(feature = "force_hash_collisions"))] + fn assert_hashes_match( + hasher: &mut CachingArrayHasher, + columns: &[ArrayRef], + random_state: &RandomState, + ) -> Result<()> { + let mut expected = vec![0; columns[0].len()]; + create_hashes(columns, random_state, &mut expected)?; + + let mut actual = vec![0; columns[0].len()]; + hasher.create_hashes(columns, &mut actual)?; + assert_eq!(actual, expected); + Ok(()) + } + + #[cfg(not(feature = "force_hash_collisions"))] + fn assert_cached_values( + hasher: &CachingArrayHasher, + column_index: usize, + expected_values: &ArrayRef, + ) -> *const u64 { + let (cached_values, value_hashes) = + hasher.dictionary_hashes(column_index).unwrap(); + assert!(Arc::ptr_eq(&cached_values, expected_values)); + value_hashes.as_ptr() + } + + #[cfg(not(feature = "force_hash_collisions"))] + #[test] + fn caching_array_hasher_hashes_regular_and_dictionary_columns() -> Result<()> { + let random_state = RandomState::with_seed(0); + let dictionary_values: ArrayRef = + Arc::new(StringArray::from(vec!["active", "inactive"])); + let first_batch = vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + string_dictionary(vec![0, 1, 0], &dictionary_values), + ]; + let mut hasher = caching_hasher(&first_batch, &random_state); + assert_hashes_match(&mut hasher, &first_batch, &random_state)?; + + assert!(hasher.dictionary_hashes(0).is_none()); + let value_hash_allocation = assert_cached_values(&hasher, 1, &dictionary_values); + + // A later batch has new plain values and dictionary keys, but shares + // the dictionary values array and therefore its cached hashes. + let second_batch = vec![ + Arc::new(Int32Array::from(vec![40, 50])) as ArrayRef, + string_dictionary(vec![1, 0], &dictionary_values), + ]; + assert_hashes_match(&mut hasher, &second_batch, &random_state)?; + assert_eq!( + assert_cached_values(&hasher, 1, &dictionary_values), + value_hash_allocation + ); + + drop(first_batch); + drop(second_batch); + drop(dictionary_values); + assert!(hasher.dictionary_hashes(1).is_none()); + + let replacement_values: ArrayRef = + Arc::new(StringArray::from(vec!["pending", "complete"])); + let third_batch = vec![ + Arc::new(Int32Array::from(vec![60, 70])) as ArrayRef, + string_dictionary(vec![0, 1], &replacement_values), + ]; + assert_hashes_match(&mut hasher, &third_batch, &random_state)?; + assert_cached_values(&hasher, 1, &replacement_values); + Ok(()) + } + + #[cfg(not(feature = "force_hash_collisions"))] + #[test] + fn caching_array_hasher_updates_dictionary_columns_independently() -> Result<()> { + let random_state = RandomState::with_seed(0); + let country_values: ArrayRef = Arc::new(StringArray::from(vec!["us", "uk"])); + let status_values: ArrayRef = + Arc::new(StringArray::from(vec!["active", "inactive"])); + let first_batch = vec![ + string_dictionary(vec![0, 1, 0], &country_values), + string_dictionary(vec![1, 0, 1], &status_values), + ]; + let mut hasher = caching_hasher(&first_batch, &random_state); + assert_hashes_match(&mut hasher, &first_batch, &random_state)?; + let country_hash_allocation = assert_cached_values(&hasher, 0, &country_values); + let status_hash_allocation = assert_cached_values(&hasher, 1, &status_values); + + // Aggregation batches can rebuild one dictionary while another group + // column continues to share its values array. + let new_country_values: ArrayRef = + Arc::new(StringArray::from(vec![Some("ca"), None])); + let second_batch = vec![ + string_dictionary(vec![1, 0], &new_country_values), + string_dictionary(vec![0, 1], &status_values), + ]; + assert_hashes_match(&mut hasher, &second_batch, &random_state)?; + assert_eq!( + assert_cached_values(&hasher, 0, &new_country_values), + country_hash_allocation + ); + assert_eq!( + assert_cached_values(&hasher, 1, &status_values), + status_hash_allocation + ); + assert_eq!(hasher.dictionary_hashes(0).unwrap().1[1], 0); + + // Reverse the transition: the first dictionary stays shared while the + // second receives a new values array. + let new_status_values: ArrayRef = + Arc::new(StringArray::from(vec!["pending", "complete"])); + let third_batch = vec![ + string_dictionary(vec![0, 1], &new_country_values), + string_dictionary(vec![1, 0], &new_status_values), + ]; + assert_hashes_match(&mut hasher, &third_batch, &random_state)?; + assert_eq!( + assert_cached_values(&hasher, 0, &new_country_values), + country_hash_allocation + ); + assert_eq!( + assert_cached_values(&hasher, 1, &new_status_values), + status_hash_allocation + ); + Ok(()) + } + + #[cfg(not(feature = "force_hash_collisions"))] + #[test] + fn caching_array_hasher_tracks_repeated_changes_with_regular_column() -> Result<()> { + let random_state = RandomState::with_seed(0); + let region_values_v1: ArrayRef = + Arc::new(StringArray::from(vec!["east", "west"])); + let category_values_v1: ArrayRef = + Arc::new(StringArray::from(vec!["hardware", "software"])); + let first_batch = vec![ + string_dictionary(vec![0, 1, 0], ®ion_values_v1), + string_dictionary(vec![1, 0, 1], &category_values_v1), + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + ]; + let mut hasher = caching_hasher(&first_batch, &random_state); + assert_hashes_match(&mut hasher, &first_batch, &random_state)?; + + // This mirrors repeated aggregation input: dictionary group columns + // can change independently while regular group columns change normally. + let region_values_v2: ArrayRef = + Arc::new(StringArray::from(vec!["north", "south"])); + let second_batch = vec![ + string_dictionary(vec![1, 0, 1], ®ion_values_v2), + string_dictionary(vec![0, 1, 0], &category_values_v1), + Arc::new(Int32Array::from(vec![40, 50, 60])) as ArrayRef, + ]; + assert_hashes_match(&mut hasher, &second_batch, &random_state)?; + + let category_values_v2: ArrayRef = + Arc::new(StringArray::from(vec!["services", "support"])); + let third_batch = vec![ + string_dictionary(vec![0, 1], ®ion_values_v2), + string_dictionary(vec![1, 0], &category_values_v2), + Arc::new(Int32Array::from(vec![70, 80])) as ArrayRef, + ]; + assert_hashes_match(&mut hasher, &third_batch, &random_state)?; + + // Replacing both dictionaries together must also update both cache + // slots without affecting the regular column. + let region_values_v3: ArrayRef = + Arc::new(StringArray::from(vec!["central", "coastal", "remote"])); + let category_values_v3: ArrayRef = Arc::new(StringArray::from(vec!["other"])); + let fourth_batch = vec![ + string_dictionary(vec![2, 0, 1], ®ion_values_v3), + string_dictionary(vec![0, 0, 0], &category_values_v3), + Arc::new(Int32Array::from(vec![90, 100, 110])) as ArrayRef, + ]; + assert_hashes_match(&mut hasher, &fourth_batch, &random_state)?; + assert_cached_values(&hasher, 0, ®ion_values_v3); + assert_cached_values(&hasher, 1, &category_values_v3); + assert!(hasher.dictionary_hashes(2).is_none()); + Ok(()) + } + #[test] fn create_hashes_for_decimal_array() -> Result<()> { let array = vec![1, 2, 3, 4] From e44de99523e8b84564295dd428d7f6cce98bcc70 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 27 Oct 2026 13:45:00 -0400 Subject: [PATCH 2/5] perf: reuse dictionary hashes in repartition --- .../physical-plan/src/repartition/mod.rs | 95 +++++++++++++++++-- 1 file changed, 86 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 929f52fd1586e..e29bd6a4953d6 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -20,6 +20,8 @@ use super::{ }; use crate::coalesce::LimitedBatchCoalescer; use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::hash_utils::CachingArrayHasher; +#[cfg(test)] use crate::hash_utils::create_hashes; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; @@ -35,7 +37,7 @@ use crate::{ use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; use arrow::compute::take_arrays; -use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; use arrow_schema::SortOptions; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; @@ -644,6 +646,7 @@ enum BatchPartitionerState { partition_reducer: StrengthReducedU64, hash_buffer: Vec, indices: Vec>, + hasher: Option, }, RoundRobin { num_partitions: usize, @@ -1037,6 +1040,7 @@ impl BatchPartitioner { partition_reducer: StrengthReducedU64::new(num_partitions as u64), hash_buffer: vec![], indices: vec![vec![]; num_partitions], + hasher: None, }, timer, }) @@ -1251,6 +1255,7 @@ impl BatchPartitioner { partition_reducer, hash_buffer, indices, + hasher, } => { // Tracking time required for distributing indexes across output partitions let timer = self.timer.timer(); @@ -1261,11 +1266,24 @@ impl BatchPartitioner { hash_buffer.clear(); hash_buffer.resize(batch.num_rows(), 0); - create_hashes( - &arrays, - REPARTITION_RANDOM_STATE.random_state(), - hash_buffer, - )?; + let hasher = hasher.get_or_insert_with(|| { + let fields = arrays + .iter() + .enumerate() + .map(|(index, array)| { + Field::new( + format!("hash_{index}"), + array.data_type().clone(), + true, + ) + }) + .collect::>(); + CachingArrayHasher::new( + &Schema::new(fields), + REPARTITION_RANDOM_STATE.random_state().clone(), + ) + }); + hasher.create_hashes(&arrays, hash_buffer)?; for v in indices.iter_mut() { v.clear(); @@ -2660,10 +2678,10 @@ mod tests { }; use arrow::array::{ - Array, ArrayRef, Decimal128Array, Int64Array, StringArray, - TimestampNanosecondArray, UInt32Array, + Array, ArrayRef, Decimal128Array, DictionaryArray, Int32Array, Int64Array, + StringArray, TimestampNanosecondArray, UInt32Array, }; - use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit}; use datafusion_common::ScalarValue; use datafusion_common::cast::{as_string_array, as_uint32_array}; use datafusion_common::config::SpillCompression; @@ -2918,6 +2936,65 @@ mod tests { Ok(()) } + #[tokio::test] + async fn hash_partition_reuses_dictionary_values_across_batches() -> Result<()> { + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let schema = Arc::new(Schema::new(vec![ + Field::new("key", dictionary_type, true), + Field::new("id", DataType::UInt32, false), + ])); + let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"])); + let make_batch = |keys: Vec>, ids: Vec| -> Result { + let dictionary: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(&values), + )?); + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![dictionary, Arc::new(UInt32Array::from(ids))], + )?) + }; + let batches = vec![ + make_batch(vec![Some(0), Some(1), Some(2), None], vec![0, 1, 2, 3])?, + make_batch(vec![Some(2), Some(0), None, Some(1)], vec![4, 5, 6, 7])?, + ]; + + let mut expected = vec![Vec::new(); 4]; + for batch in &batches { + let mut hashes = vec![0; batch.num_rows()]; + create_hashes( + [batch.column(0)], + REPARTITION_RANDOM_STATE.random_state(), + &mut hashes, + )?; + let ids = as_uint32_array(batch.column(1))?; + for (hash, id) in hashes.into_iter().zip(ids.values().iter().copied()) { + expected[(hash % 4) as usize].push(id); + } + } + + let output = repartition( + &schema, + vec![batches], + Partitioning::Hash(vec![col("key", &schema)?], 4), + ) + .await?; + let mut actual = vec![Vec::new(); 4]; + for (partition, batches) in output.iter().enumerate() { + for batch in batches { + actual[partition] + .extend(as_uint32_array(batch.column(1))?.values().iter().copied()); + } + } + for values in expected.iter_mut().chain(actual.iter_mut()) { + values.sort_unstable(); + } + + assert_eq!(actual, expected); + Ok(()) + } + #[tokio::test] async fn many_to_many_range_partition() -> Result<()> { let schema = test_schema(false); From 4849641ca6dbf63f4e6b89c6663b91c2f8fa2f49 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 27 Oct 2026 15:35:18 -0400 Subject: [PATCH 3/5] perf: reuse dictionary hashes in aggregation --- .../group_values/multi_group_by/dictionary.rs | 21 ++++- .../group_values/multi_group_by/mod.rs | 88 ++++++++++++++++--- 2 files changed, 95 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index 3726c3ea2fb23..7050b8775b01e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -182,16 +182,29 @@ impl DictionaryGroupValuesColumn { /// array; without this the per-batch cost is O(dictionary cardinality) /// rather than O(rows). `take_n` remaps `inner` and clears /// `cached_values`, forcing a rebuild. - fn sync_value_cache(&mut self, dict_values: &ArrayRef) { + fn sync_value_cache_with_hashes( + &mut self, + dict_values: &ArrayRef, + hashes: Option<&[u64]>, + ) { if self.is_cached(dict_values) { return; } - self.hash_values(dict_values); + if let Some(hashes) = hashes.filter(|hashes| hashes.len() == dict_values.len()) { + self.val_hashes.clear(); + self.val_hashes.extend_from_slice(hashes); + } else { + self.hash_values(dict_values); + } self.val_to_inner.clear(); self.val_to_inner.resize(dict_values.len(), usize::MAX); self.cached_values = Some(Arc::downgrade(dict_values)); } + fn sync_value_cache(&mut self, dict_values: &ArrayRef) { + self.sync_value_cache_with_hashes(dict_values, None); + } + fn find_or_insert_value( &mut self, dict_values: &ArrayRef, @@ -338,6 +351,10 @@ impl GroupColumn self.check_key_overflow() } + fn set_dictionary_hashes(&mut self, values: &ArrayRef, hashes: &[u64]) { + self.sync_value_cache_with_hashes(values, Some(hashes)); + } + fn vectorized_equal_to( &mut self, lhs_rows: &[usize], 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 bf784ea5d1591..82ee8823c83ac 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 @@ -25,16 +25,13 @@ use crate::aggregates::group_values::multi_group_by::{ fixed_size_binary::FixedSizeBinaryGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, row_backed::RowsGroupColumn, }; -use arrow::array::{ - Array, ArrayRef, BooleanBufferBuilder, downcast_primitive, new_empty_array, -}; +use arrow::array::{ArrayRef, BooleanBufferBuilder, downcast_primitive, new_empty_array}; use arrow::datatypes::{ BinaryViewType, DataType, Field, Int8Type, Int16Type, Int32Type, Int64Type, Schema, SchemaRef, StringViewType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; 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::hash_utils::CachingArrayHasher; use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::{EmitTo, GroupSelection}; @@ -82,6 +79,9 @@ pub trait GroupColumn: Send + Sync { /// The vectorized version `append_val` fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()>; + /// Installs hashes already computed for a dictionary values array. + fn set_dictionary_hashes(&mut self, _values: &ArrayRef, _hashes: &[u64]) {} + /// Returns the number of rows stored in this builder fn len(&self) -> usize; @@ -216,8 +216,8 @@ pub struct GroupValuesColumn { /// reused buffer to store hashes hashes_buffer: Vec, - /// Random state for creating hashes - random_state: RandomState, + /// Caches dictionary value hashes by group column. + hasher: CachingArrayHasher, } /// Buffers to store intermediate results in `vectorized_append` @@ -288,6 +288,10 @@ impl GroupValuesColumn { pub fn try_new(schema: SchemaRef) -> Result { let map = HashTable::with_capacity(0); let group_values = Self::build_group_columns(&schema)?; + let hasher = CachingArrayHasher::new( + schema.as_ref(), + crate::aggregates::AGGREGATION_HASH_SEED, + ); Ok(Self { schema, map, @@ -297,10 +301,25 @@ impl GroupValuesColumn { map_size: 0, group_values, hashes_buffer: Default::default(), - random_state: crate::aggregates::AGGREGATION_HASH_SEED, + hasher, }) } + fn create_hashes_with_cache( + hasher: &mut CachingArrayHasher, + group_values: &mut [Box], + cols: &[ArrayRef], + hashes: &mut [u64], + ) -> Result<()> { + hasher.create_hashes(cols, hashes)?; + for (index, group_value) in group_values.iter_mut().enumerate() { + if let Some((values, value_hashes)) = hasher.dictionary_hashes(index) { + group_value.set_dictionary_hashes(&values, value_hashes); + } + } + Ok(()) + } + /// Build one fresh [`GroupColumn`] per field in the schema. /// /// Used at construction time (`try_new`) and to repopulate the column @@ -373,7 +392,12 @@ impl GroupValuesColumn { let batch_hashes = &mut self.hashes_buffer; batch_hashes.clear(); batch_hashes.resize(n_rows, 0); - create_hashes(cols, &self.random_state, batch_hashes)?; + Self::create_hashes_with_cache( + &mut self.hasher, + &mut self.group_values, + cols, + batch_hashes, + )?; for (row, &target_hash) in batch_hashes.iter().enumerate() { let entry = self @@ -474,7 +498,12 @@ impl GroupValuesColumn { let mut batch_hashes = mem::take(&mut self.hashes_buffer); batch_hashes.clear(); batch_hashes.resize(n_rows, 0); - create_hashes(cols, &self.random_state, &mut batch_hashes)?; + Self::create_hashes_with_cache( + &mut self.hasher, + &mut self.group_values, + cols, + &mut batch_hashes, + )?; // General steps for one round `vectorized equal_to & append`: // 1. Collect vectorized context by checking hash values of `cols` in `map`, @@ -1090,6 +1119,7 @@ impl GroupValues for GroupValuesColumn { + self.emit_group_index_list_buffer.allocated_size() + self.vectorized_operation_buffers.size() + self.hashes_buffer.allocated_size() + + self.hasher.size() } fn is_empty(&self) -> bool { @@ -1247,6 +1277,10 @@ impl GroupValues for GroupValuesColumn { self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); self.hashes_buffer.clear(); self.hashes_buffer.shrink_to(num_rows); + self.hasher = CachingArrayHasher::new( + self.schema.as_ref(), + crate::aggregates::AGGREGATION_HASH_SEED, + ); // Such structures are only used in `non-streaming` case if !STREAMING { @@ -2106,15 +2140,45 @@ mod tests { Field::new( "category", DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), - false, + true, ), ])); let mut gv = GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); gv.intern(&[labels, categories], &mut vec![]).unwrap(); + let (cached, _) = gv.hasher.dictionary_hashes(1).unwrap(); + assert!(Arc::ptr_eq(&cached, &dict_vocab)); + let batch = |values| -> Vec { + vec![ + Arc::new(StringArray::from(vec!["g0"])), + Arc::new(DictionaryArray::::new( + Int8Array::from(vec![0]), + values, + )), + ] + }; + gv.intern(&batch(Arc::clone(&dict_vocab)), &mut vec![]) + .unwrap(); + let (cached, _) = gv.hasher.dictionary_hashes(1).unwrap(); + assert!(Arc::ptr_eq(&cached, &dict_vocab)); + + let new_values: ArrayRef = + Arc::new(StringArray::from(vec![Some("cat"), None, Some("dog")])); + let null_batch: Vec = vec![ + Arc::new(StringArray::from(vec!["null", "null"])), + Arc::new(DictionaryArray::::new( + Int8Array::from(vec![None, Some(1)]), + Arc::clone(&new_values), + )), + ]; + let mut null_groups = vec![]; + gv.intern(&null_batch, &mut null_groups).unwrap(); + let (cached, _) = gv.hasher.dictionary_hashes(1).unwrap(); + assert!(Arc::ptr_eq(&cached, &new_values)); + assert_eq!(null_groups[0], null_groups[1]); let out = gv.emit(EmitTo::All).unwrap(); - assert_eq!(out[0].len(), n_groups); + assert_eq!(out[0].len(), n_groups + 1); assert!(matches!( out[1].data_type(), DataType::Dictionary(k, v) From a785df1315e51e6f1db1d8b4101d0befe1049a96 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Fri, 30 Oct 2026 17:02:52 -0400 Subject: [PATCH 4/5] perf(aggregate): reuse dictionary group mappings For a single non-streaming dictionary group key, the outer table already guarantees that appended values are distinct groups. Reuse the dictionary column's cached value-to-slot mapping as group indices so later batches can skip hashing and probing. --- .../group_values/multi_group_by/dictionary.rs | 104 ++++++++++++++---- .../group_values/multi_group_by/mod.rs | 100 ++++++++++++++++- 2 files changed, 183 insertions(+), 21 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index 7050b8775b01e..67d98e644dacb 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -54,6 +54,9 @@ pub struct DictionaryGroupValuesColumn /// were built for. A matching live array reuses both caches without /// retaining an input dictionary after its batch is released. cached_values: Option>, + /// Whether the outer table guarantees each append creates a new group, so + /// inner slots and group indices are identical. + values_are_group_unique: bool, _phantom: PhantomData, } @@ -71,6 +74,7 @@ impl DictionaryGroupValuesColumn { val_to_inner: Vec::default(), val_hashes: Vec::default(), cached_values: None, + values_are_group_unique: false, _phantom: PhantomData, } } @@ -190,6 +194,13 @@ impl DictionaryGroupValuesColumn { if self.is_cached(dict_values) { return; } + if self.values_are_group_unique { + self.val_hashes.clear(); + self.val_to_inner.clear(); + self.val_to_inner.resize(dict_values.len(), usize::MAX); + self.cached_values = Some(Arc::downgrade(dict_values)); + return; + } if let Some(hashes) = hashes.filter(|hashes| hashes.len() == dict_values.len()) { self.val_hashes.clear(); self.val_hashes.extend_from_slice(hashes); @@ -205,6 +216,23 @@ impl DictionaryGroupValuesColumn { self.sync_value_cache_with_hashes(dict_values, None); } + fn insert_value(&mut self, dict_values: &ArrayRef, val_idx: usize) -> Result { + if self.values_are_group_unique { + let slot = self.inner.len(); + self.inner.append_val(dict_values, val_idx)?; + Ok(slot) + } else { + self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx]) + } + } + + fn push_inner_slot(&mut self, inner_slot: usize) { + debug_assert!( + !self.values_are_group_unique || inner_slot == self.group_to_inner.len() + ); + self.group_to_inner.push(inner_slot); + } + fn find_or_insert_value( &mut self, dict_values: &ArrayRef, @@ -344,10 +372,10 @@ impl GroupColumn Some(val_idx) => { let dict_values = dict.values(); self.sync_value_cache(dict_values); - self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx])? + self.insert_value(dict_values, val_idx)? } }; - self.group_to_inner.push(inner_slot); + self.push_inner_slot(inner_slot); self.check_key_overflow() } @@ -355,6 +383,43 @@ impl GroupColumn self.sync_value_cache_with_hashes(values, Some(hashes)); } + fn try_map_dictionary_groups( + &self, + array: &ArrayRef, + groups: &mut Vec, + ) -> bool { + if !self.values_are_group_unique { + return false; + } + + let dictionary = array.as_dictionary::(); + if !self.is_cached(dictionary.values()) { + return false; + } + + let keys = dictionary.keys(); + let raw_keys = keys.values(); + let nulls = keys.nulls(); + groups.clear(); + for row in 0..dictionary.len() { + let group = if nulls.is_some_and(|nulls| nulls.is_null(row)) { + self.null_inner_slot.unwrap_or(usize::MAX) + } else { + self.val_to_inner[raw_keys[row].as_usize()] + }; + if group == usize::MAX { + groups.clear(); + return false; + } + groups.push(group); + } + true + } + + fn set_values_are_group_unique(&mut self) { + self.values_are_group_unique = true; + } + fn vectorized_equal_to( &mut self, lhs_rows: &[usize], @@ -469,14 +534,10 @@ impl GroupColumn self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { self.find_or_insert_null()? } else { - self.find_or_insert_value( - dict_values, - val_idx, - self.val_hashes[val_idx], - )? + self.insert_value(dict_values, val_idx)? }; } - self.group_to_inner.push(self.val_to_inner[val_idx]); + self.push_inner_slot(self.val_to_inner[val_idx]); } } else { let null_buf = dict_keys.nulls().unwrap(); @@ -489,16 +550,12 @@ impl GroupColumn self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { self.find_or_insert_null()? } else { - self.find_or_insert_value( - dict_values, - val_idx, - self.val_hashes[val_idx], - )? + self.insert_value(dict_values, val_idx)? }; } self.val_to_inner[val_idx] }; - self.group_to_inner.push(slot); + self.push_inner_slot(slot); } } @@ -615,7 +672,12 @@ impl GroupColumn self.value_dedup_size = 0; self.null_inner_slot = None; - self.hash_values(&all_inner_values); + if self.values_are_group_unique { + self.cached_values = None; + self.val_hashes.clear(); + } else { + self.hash_values(&all_inner_values); + } for (new_slot, &old_slot) in new_to_old.iter().enumerate() { if all_inner_values.is_null(old_slot) { @@ -627,11 +689,13 @@ impl GroupColumn self.inner .append_val(&all_inner_values, old_slot) .expect("append value failed in take_n"); - self.value_dedup.insert_accounted( - (self.val_hashes[old_slot], new_slot), - |&(entry_hash, _)| entry_hash, - &mut self.value_dedup_size, - ); + if !self.values_are_group_unique { + self.value_dedup.insert_accounted( + (self.val_hashes[old_slot], new_slot), + |&(entry_hash, _)| entry_hash, + &mut self.value_dedup_size, + ); + } } } 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 82ee8823c83ac..183cf6b1dde84 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 @@ -82,6 +82,19 @@ pub trait GroupColumn: Send + Sync { /// Installs hashes already computed for a dictionary values array. fn set_dictionary_hashes(&mut self, _values: &ArrayRef, _hashes: &[u64]) {} + /// Maps dictionary keys directly to group indices when all referenced + /// values are already cached. + fn try_map_dictionary_groups( + &self, + _array: &ArrayRef, + _groups: &mut Vec, + ) -> bool { + false + } + + /// Marks values appended to this column as distinct groups. + fn set_values_are_group_unique(&mut self) {} + /// Returns the number of rows stored in this builder fn len(&self) -> usize; @@ -218,6 +231,9 @@ pub struct GroupValuesColumn { /// Caches dictionary value hashes by group column. hasher: CachingArrayHasher, + + /// Whether the non-streaming single dictionary-key fast path applies. + map_dictionary_groups: bool, } /// Buffers to store intermediate results in `vectorized_append` @@ -292,6 +308,9 @@ impl GroupValuesColumn { schema.as_ref(), crate::aggregates::AGGREGATION_HASH_SEED, ); + let map_dictionary_groups = !STREAMING + && schema.fields().len() == 1 + && matches!(schema.field(0).data_type(), DataType::Dictionary(_, _)); Ok(Self { schema, map, @@ -302,6 +321,7 @@ impl GroupValuesColumn { group_values, hashes_buffer: Default::default(), hasher, + map_dictionary_groups, }) } @@ -332,6 +352,9 @@ impl GroupValuesColumn { for f in schema.fields().iter() { v.push(make_group_column(f.as_ref())?); } + if !STREAMING && v.len() == 1 { + v[0].set_values_are_group_unique(); + } Ok(v) } @@ -1092,6 +1115,11 @@ impl GroupValues for GroupValuesColumn { // `try_new` and the reset points in `emit` / `clear_shrink` keep // `self.group_values` populated with one builder per schema field, // so no lazy initialization is needed here. + if self.map_dictionary_groups + && self.group_values[0].try_map_dictionary_groups(&cols[0], groups) + { + return Ok(()); + } if !STREAMING { self.vectorized_intern(cols, groups) } else { @@ -1281,7 +1309,6 @@ impl GroupValues for GroupValuesColumn { self.schema.as_ref(), crate::aggregates::AGGREGATION_HASH_SEED, ); - // Such structures are only used in `non-streaming` case if !STREAMING { self.group_index_lists.clear(); @@ -2188,6 +2215,77 @@ mod tests { assert_eq!(out[1].as_dictionary::().values().len(), 2); } + #[test] + fn single_dictionary_group_cache_tracks_values_and_emits() { + use arrow::array::{DictionaryArray, Int8Array}; + use arrow::datatypes::Int8Type; + + let schema = Arc::new(Schema::new(vec![Field::new( + "category", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + )])); + let mut group_values = GroupValuesColumn::::try_new(schema).unwrap(); + let first_values: ArrayRef = + Arc::new(StringArray::from(vec![Some("cat"), Some("dog"), None])); + let dictionary = |keys, values| -> ArrayRef { + Arc::new(DictionaryArray::::new( + Int8Array::from(keys), + values, + )) + }; + + let mut groups = vec![]; + group_values + .intern( + &[dictionary( + vec![Some(0), Some(1), None, Some(2)], + Arc::clone(&first_values), + )], + &mut groups, + ) + .unwrap(); + assert_eq!(groups, vec![0, 1, 2, 2]); + + // Every value index is now resolved, so another batch sharing the + // values array can map keys directly to groups. + let cached_batch = dictionary( + vec![Some(1), Some(0), Some(2), None], + Arc::clone(&first_values), + ); + assert!( + group_values.group_values[0] + .try_map_dictionary_groups(&cached_batch, &mut groups) + ); + assert_eq!(groups, vec![1, 0, 2, 2]); + group_values.intern(&[cached_batch], &mut groups).unwrap(); + assert_eq!(groups, vec![1, 0, 2, 2]); + + // A new values array uses normal logical-value matching to rebuild the + // cache even when its key assignments differ. + let reordered_values: ArrayRef = + Arc::new(StringArray::from(vec![Some("dog"), Some("cat"), None])); + group_values + .intern( + &[dictionary( + vec![Some(0), Some(1), Some(2)], + Arc::clone(&reordered_values), + )], + &mut groups, + ) + .unwrap(); + assert_eq!(groups, vec![1, 0, 2]); + + group_values.emit(EmitTo::First(1)).unwrap(); + let after_emit = dictionary(vec![Some(0), Some(2), Some(1)], reordered_values); + assert!( + !group_values.group_values[0] + .try_map_dictionary_groups(&after_emit, &mut groups) + ); + group_values.intern(&[after_emit], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2]); + } + #[test] fn test_intern_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); From f02180725bc518d0b8bce2af28706bd3a66009f1 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Fri, 30 Oct 2026 19:29:28 -0400 Subject: [PATCH 5/5] perf(aggregate): limit dictionary cache plumbing Use the caching hasher only for non-streaming group schemas that contain dictionary columns, and share cached hashes only with those columns. Streaming and non-dictionary aggregates keep the existing hashing path without per-batch cache plumbing. --- .../group_values/multi_group_by/mod.rs | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) 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 183cf6b1dde84..7b58f85b0fc92 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 @@ -31,7 +31,7 @@ use arrow::datatypes::{ SchemaRef, StringViewType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use datafusion_common::float_semantics::{has_float_leaf, normalize_floats}; -use datafusion_common::hash_utils::CachingArrayHasher; +use datafusion_common::hash_utils::{CachingArrayHasher, create_hashes}; use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::{EmitTo, GroupSelection}; @@ -232,6 +232,9 @@ pub struct GroupValuesColumn { /// Caches dictionary value hashes by group column. hasher: CachingArrayHasher, + /// Group columns whose dictionary hashes need to be shared with storage. + dictionary_indices: Vec, + /// Whether the non-streaming single dictionary-key fast path applies. map_dictionary_groups: bool, } @@ -304,6 +307,19 @@ impl GroupValuesColumn { pub fn try_new(schema: SchemaRef) -> Result { let map = HashTable::with_capacity(0); let group_values = Self::build_group_columns(&schema)?; + let dictionary_indices = if STREAMING { + vec![] + } else { + schema + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| { + matches!(field.data_type(), DataType::Dictionary(_, _)) + .then_some(index) + }) + .collect() + }; let hasher = CachingArrayHasher::new( schema.as_ref(), crate::aggregates::AGGREGATION_HASH_SEED, @@ -321,6 +337,7 @@ impl GroupValuesColumn { group_values, hashes_buffer: Default::default(), hasher, + dictionary_indices, map_dictionary_groups, }) } @@ -328,11 +345,18 @@ impl GroupValuesColumn { fn create_hashes_with_cache( hasher: &mut CachingArrayHasher, group_values: &mut [Box], + dictionary_indices: &[usize], cols: &[ArrayRef], hashes: &mut [u64], ) -> Result<()> { + if dictionary_indices.is_empty() { + create_hashes(cols, &crate::aggregates::AGGREGATION_HASH_SEED, hashes)?; + return Ok(()); + } + hasher.create_hashes(cols, hashes)?; - for (index, group_value) in group_values.iter_mut().enumerate() { + for &index in dictionary_indices { + let group_value = &mut group_values[index]; if let Some((values, value_hashes)) = hasher.dictionary_hashes(index) { group_value.set_dictionary_hashes(&values, value_hashes); } @@ -418,6 +442,7 @@ impl GroupValuesColumn { Self::create_hashes_with_cache( &mut self.hasher, &mut self.group_values, + &self.dictionary_indices, cols, batch_hashes, )?; @@ -524,6 +549,7 @@ impl GroupValuesColumn { Self::create_hashes_with_cache( &mut self.hasher, &mut self.group_values, + &self.dictionary_indices, cols, &mut batch_hashes, )?; @@ -1147,6 +1173,7 @@ impl GroupValues for GroupValuesColumn { + self.emit_group_index_list_buffer.allocated_size() + self.vectorized_operation_buffers.size() + self.hashes_buffer.allocated_size() + + self.dictionary_indices.allocated_size() + self.hasher.size() }