From 74f019e1e71e64af0cf582bdef452502a36fda1a Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 27 Oct 2026 14:34:44 -0400 Subject: [PATCH 1/4] refactor: add reusable caching array hasher Reuse dictionary value hashes while the values Arc is unchanged so operators do not rehash the same dictionary for every batch. --- datafusion/common/src/hash_utils.rs | 172 ++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index c9cc48a11dd21..37166e21ae03a 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,101 @@ use std::hash::{BuildHasher, Hash, Hasher}; pub type RandomState = FixedState; pub type QualityRandomState = foldhash::quality::FixedState; +type DictionaryHashEntry = (Weak, Vec); + +/// Hashes one array at a time, reusing dictionary value hashes when possible. +#[derive(Debug)] +pub struct CachingArrayHasher { + /// Stored with the cache so value hashes cannot cross hash seeds. + random_state: RandomState, + /// Weak identity avoids retaining values; the hash allocation is reusable. + dictionary: Option, +} + +impl CachingArrayHasher { + /// Creates an empty cache for `random_state`. + pub fn new(random_state: RandomState) -> Self { + Self { + random_state, + dictionary: None, + } + } + + /// Hashes `array`, reusing dictionary value hashes while its values array + /// is unchanged. + /// + /// `hashes_buffer` must have the same length as `array`. If `rehash` is + /// `false`, this initializes each row hash; if `true`, it combines the + /// array hash with the existing row hash. + pub fn create_hashes<'a>( + &mut self, + array: &dyn Array, + hashes_buffer: &'a mut [u64], + rehash: bool, + ) -> Result<&'a mut [u64]> { + #[cfg(feature = "force_hash_collisions")] + { + hash_single_array(array, &self.random_state, hashes_buffer, rehash)?; + } + + #[cfg(not(feature = "force_hash_collisions"))] + { + if !matches!(array.data_type(), DataType::Dictionary(_, _)) { + self.dictionary = None; + hash_single_array(array, &self.random_state, hashes_buffer, rehash)?; + return Ok(hashes_buffer); + } + + downcast_dictionary_array! { + array => { + let values = array.values(); + if !self.dictionary.as_ref().is_some_and(|(cached, _)| { + std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values)) + }) { + let mut hashes = self + .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, + )?; + self.dictionary = Some((Arc::downgrade(values), hashes)); + } + dispatch_dictionary_scatter( + array, + &self.dictionary.as_ref().unwrap().1, + hashes_buffer, + rehash, + ); + } + _ => unreachable!(), + }; + } + + Ok(hashes_buffer) + } + + /// Returns the current dictionary values and hashes, if any. + pub fn dictionary_hashes(&self) -> Option<(ArrayRef, &[u64])> { + let (values, hashes) = self.dictionary.as_ref()?; + Some((values.upgrade()?, hashes)) + } + + /// Returns the allocated size of the cache. + pub fn size(&self) -> usize { + self.dictionary + .as_ref() + .map_or(0, |entry| entry.1.capacity() * size_of::()) + } +} + /// Fixed quality hash state used by HyperLogLog sketches. /// /// The seed is part of the HLL wire/storage semantics: serialized sketches only @@ -1357,6 +1456,79 @@ mod tests { } } + #[cfg(not(feature = "force_hash_collisions"))] + #[test] + fn caching_array_hasher_reuses_dictionary_allocation() -> Result<()> { + let random_state = RandomState::with_seed(0); + let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let dictionary = DictionaryArray::::new( + Int8Array::from(vec![0, 1, 0]), + Arc::clone(&values), + ); + let mut expected = vec![0; dictionary.len()]; + create_hashes([&dictionary as &dyn Array], &random_state, &mut expected)?; + + let mut cache = CachingArrayHasher::new(random_state.clone()); + let mut actual = vec![0; dictionary.len()]; + cache.create_hashes(&dictionary, &mut actual, false)?; + assert_eq!(actual, expected); + + let (cached_values, value_hashes) = cache.dictionary_hashes().unwrap(); + assert!(Arc::ptr_eq(&cached_values, &values)); + let allocation = value_hashes.as_ptr(); + + let prefix = Int32Array::from(vec![10, 20]); + let dictionary = DictionaryArray::::new( + Int8Array::from(vec![1, 0]), + Arc::clone(&values), + ); + let mut expected = vec![0; dictionary.len()]; + create_hashes( + [&prefix as &dyn Array, &dictionary], + &random_state, + &mut expected, + )?; + let mut actual = vec![0; dictionary.len()]; + create_hashes([&prefix as &dyn Array], &random_state, &mut actual)?; + cache.create_hashes(&dictionary, &mut actual, true)?; + assert_eq!(actual, expected); + + let new_values: ArrayRef = Arc::new(StringArray::from(vec![Some("c"), None])); + let dictionary = DictionaryArray::::new( + Int8Array::from(vec![1, 0]), + Arc::clone(&new_values), + ); + actual.clear(); + actual.resize(dictionary.len(), 0); + expected.clear(); + expected.resize(dictionary.len(), 0); + create_hashes([&dictionary as &dyn Array], &random_state, &mut expected)?; + cache.create_hashes(&dictionary, &mut actual, false)?; + + let (cached_values, value_hashes) = cache.dictionary_hashes().unwrap(); + assert_eq!(actual, expected); + assert!(Arc::ptr_eq(&cached_values, &new_values)); + assert_eq!(value_hashes.as_ptr(), allocation); + assert_eq!(value_hashes[1], 0); + Ok(()) + } + + #[test] + fn caching_array_hasher_hashes_non_dictionary() -> Result<()> { + let random_state = RandomState::with_seed(0); + let array = Int32Array::from(vec![1, 2, 3]); + let mut expected = vec![0; array.len()]; + create_hashes([&array as &dyn Array], &random_state, &mut expected)?; + + let mut cache = CachingArrayHasher::new(random_state); + let mut actual = vec![0; array.len()]; + cache.create_hashes(&array, &mut actual, false)?; + + assert_eq!(actual, expected); + assert!(cache.dictionary_hashes().is_none()); + Ok(()) + } + #[test] fn create_hashes_for_decimal_array() -> Result<()> { let array = vec![1, 2, 3, 4] From 049070072b98b6150032ccb14bbd2cbf529a9ec9 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Wed, 28 Oct 2026 00:34:43 -0400 Subject: [PATCH 2/4] test: consolidate caching array hasher coverage --- datafusion/common/src/hash_utils.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index 37166e21ae03a..08e85601fed0a 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -1458,7 +1458,7 @@ mod tests { #[cfg(not(feature = "force_hash_collisions"))] #[test] - fn caching_array_hasher_reuses_dictionary_allocation() -> Result<()> { + fn caching_array_hasher() -> Result<()> { let random_state = RandomState::with_seed(0); let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); let dictionary = DictionaryArray::::new( @@ -1510,17 +1510,11 @@ mod tests { assert!(Arc::ptr_eq(&cached_values, &new_values)); assert_eq!(value_hashes.as_ptr(), allocation); assert_eq!(value_hashes[1], 0); - Ok(()) - } - #[test] - fn caching_array_hasher_hashes_non_dictionary() -> Result<()> { - let random_state = RandomState::with_seed(0); let array = Int32Array::from(vec![1, 2, 3]); let mut expected = vec![0; array.len()]; create_hashes([&array as &dyn Array], &random_state, &mut expected)?; - let mut cache = CachingArrayHasher::new(random_state); let mut actual = vec![0; array.len()]; cache.create_hashes(&array, &mut actual, false)?; From 9c32dbaaf42c30282abb4be294ea69aa9f53aa4d Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Wed, 28 Oct 2026 14:30:46 -0400 Subject: [PATCH 3/4] refactor: cache dictionary hashes by column --- datafusion/common/src/hash_utils.rs | 322 ++++++++++++++++++++-------- 1 file changed, 238 insertions(+), 84 deletions(-) diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index 08e85601fed0a..419b0ee0fc958 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -34,57 +34,86 @@ pub type QualityRandomState = foldhash::quality::FixedState; type DictionaryHashEntry = (Weak, Vec); -/// Hashes one array at a time, reusing dictionary value hashes when possible. +/// 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, - /// Weak identity avoids retaining values; the hash allocation is reusable. - dictionary: Option, + /// One entry per column. Weak identity avoids retaining dictionary values. + dictionaries: Vec>, } impl CachingArrayHasher { - /// Creates an empty cache for `random_state`. - pub fn new(random_state: RandomState) -> Self { + /// 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, - dictionary: None, + dictionaries: schema.fields().iter().map(|_| None).collect(), } } - /// Hashes `array`, reusing dictionary value hashes while its values array - /// is unchanged. - /// - /// `hashes_buffer` must have the same length as `array`. If `rehash` is - /// `false`, this initializes each row hash; if `true`, it combines the - /// array hash with the existing row hash. - pub fn create_hashes<'a>( + /// Hashes `arrays`, reusing dictionary value hashes while each column's + /// values array is unchanged. + pub fn create_hashes<'a, I, T>( &mut self, - array: &dyn Array, + 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<&'a mut [u64]> { + ) -> 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(_, _)) { - self.dictionary = None; + *dictionary = None; hash_single_array(array, &self.random_state, hashes_buffer, rehash)?; - return Ok(hashes_buffer); + return Ok(()); } downcast_dictionary_array! { array => { let values = array.values(); - if !self.dictionary.as_ref().is_some_and(|(cached, _)| { + if !dictionary.as_ref().is_some_and(|(cached, _)| { std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values)) }) { - let mut hashes = self - .dictionary + let mut hashes = dictionary .take() .map(|(_, hashes)| hashes) .unwrap_or_default(); @@ -97,11 +126,11 @@ impl CachingArrayHasher { &mut hashes, false, )?; - self.dictionary = Some((Arc::downgrade(values), hashes)); + *dictionary = Some((Arc::downgrade(values), hashes)); } dispatch_dictionary_scatter( array, - &self.dictionary.as_ref().unwrap().1, + &dictionary.as_ref().unwrap().1, hashes_buffer, rehash, ); @@ -110,20 +139,24 @@ impl CachingArrayHasher { }; } - Ok(hashes_buffer) + Ok(()) } - /// Returns the current dictionary values and hashes, if any. - pub fn dictionary_hashes(&self) -> Option<(ArrayRef, &[u64])> { - let (values, hashes) = self.dictionary.as_ref()?; + /// Returns the current dictionary values and hashes for `index`, if any. + 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.dictionary - .as_ref() - .map_or(0, |entry| entry.1.capacity() * size_of::()) + self.dictionaries.capacity() * size_of::>() + + self + .dictionaries + .iter() + .flatten() + .map(|entry| entry.1.capacity() * size_of::()) + .sum::() } } @@ -1457,69 +1490,190 @@ mod tests { } #[cfg(not(feature = "force_hash_collisions"))] - #[test] - fn caching_array_hasher() -> Result<()> { - let random_state = RandomState::with_seed(0); - let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); - let dictionary = DictionaryArray::::new( - Int8Array::from(vec![0, 1, 0]), - Arc::clone(&values), - ); - let mut expected = vec![0; dictionary.len()]; - create_hashes([&dictionary as &dyn Array], &random_state, &mut expected)?; + fn string_dictionary(keys: Vec, values: &ArrayRef) -> ArrayRef { + Arc::new(DictionaryArray::::new( + Int8Array::from(keys), + Arc::clone(values), + )) + } - let mut cache = CachingArrayHasher::new(random_state.clone()); - let mut actual = vec![0; dictionary.len()]; - cache.create_hashes(&dictionary, &mut actual, false)?; + #[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(()) + } - let (cached_values, value_hashes) = cache.dictionary_hashes().unwrap(); - assert!(Arc::ptr_eq(&cached_values, &values)); - let allocation = value_hashes.as_ptr(); + #[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() + } - let prefix = Int32Array::from(vec![10, 20]); - let dictionary = DictionaryArray::::new( - Int8Array::from(vec![1, 0]), - Arc::clone(&values), - ); - let mut expected = vec![0; dictionary.len()]; - create_hashes( - [&prefix as &dyn Array, &dictionary], - &random_state, - &mut expected, - )?; - let mut actual = vec![0; dictionary.len()]; - create_hashes([&prefix as &dyn Array], &random_state, &mut actual)?; - cache.create_hashes(&dictionary, &mut actual, true)?; - assert_eq!(actual, expected); + #[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)?; - let new_values: ArrayRef = Arc::new(StringArray::from(vec![Some("c"), None])); - let dictionary = DictionaryArray::::new( - Int8Array::from(vec![1, 0]), - Arc::clone(&new_values), - ); - actual.clear(); - actual.resize(dictionary.len(), 0); - expected.clear(); - expected.resize(dictionary.len(), 0); - create_hashes([&dictionary as &dyn Array], &random_state, &mut expected)?; - cache.create_hashes(&dictionary, &mut actual, false)?; - - let (cached_values, value_hashes) = cache.dictionary_hashes().unwrap(); - assert_eq!(actual, expected); - assert!(Arc::ptr_eq(&cached_values, &new_values)); - assert_eq!(value_hashes.as_ptr(), allocation); - assert_eq!(value_hashes[1], 0); + assert!(hasher.dictionary_hashes(0).is_none()); + let value_hash_allocation = assert_cached_values(&hasher, 1, &dictionary_values); - let array = Int32Array::from(vec![1, 2, 3]); - let mut expected = vec![0; array.len()]; - create_hashes([&array as &dyn Array], &random_state, &mut expected)?; + // 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 + ); + Ok(()) + } - let mut actual = vec![0; array.len()]; - cache.create_hashes(&array, &mut actual, false)?; + #[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(()) + } - assert_eq!(actual, expected); - assert!(cache.dictionary_hashes().is_none()); + #[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(()) } From d7207e94ee72eae8d990bfdab1db06bb967b312c Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 27 Oct 2026 15:51:22 -0400 Subject: [PATCH 4/4] perf: reuse dictionary value hashes in hash joins --- .../physical-plan/src/joins/hash_join/exec.rs | 17 ++++++--- .../src/joins/hash_join/stream.rs | 24 +++++++------ datafusion/physical-plan/src/joins/utils.rs | 35 +++++++++++++++++-- 3 files changed, 59 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index ddacf5a1b37a4..3c28666e1f5fe 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -31,7 +31,8 @@ use crate::joins::key_range_bitmap::KeyRangeBitmap; use crate::joins::probe_completion::{Completion, ProbeCompletion}; use crate::joins::utils::{ OnceAsync, OnceFut, asymmetric_join_output_partitioning, emits_unmatched_left_rows, - is_existence_join, reorder_output_after_swap, swap_join_projection, update_hash, + is_existence_join, reorder_output_after_swap, swap_join_projection, + update_hash_with_cache, }; use crate::joins::{JoinOn, JoinOnRef, PartitionMode, SharedBitmapBuilder}; use crate::metrics::{Count, MetricBuilder, MetricCategory}; @@ -86,7 +87,7 @@ use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, l use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; -use datafusion_common::hash_utils::RandomState; +use datafusion_common::hash_utils::{CachingArrayHasher, RandomState}; use datafusion_physical_expr_common::physical_expr::fmt_sql; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::TryStreamExt; @@ -1796,7 +1797,7 @@ impl ExecutionPlan for HashJoinExec { build_accumulator, self.mode, self.fetch, - ))) + )?)) } fn metrics(&self) -> Option { @@ -3234,6 +3235,12 @@ fn build_in_memory( } // The maximum is known: avoid geometric growth and its excess capacity. let mut hashes_buffer = Vec::with_capacity(max_batch_rows); + let fields = on_left + .iter() + .map(|expr| expr.return_field(schema)) + .collect::>>()?; + let mut hasher = + CachingArrayHasher::new(&Schema::new(fields), random_state.clone()); let mut offset = num_rows; // Update the hashmap starting from the last batch, so rows with smaller @@ -3243,16 +3250,16 @@ fn build_in_memory( offset -= batch.num_rows(); hashes_buffer.clear(); hashes_buffer.resize(batch.num_rows(), 0); - update_hash( + update_hash_with_cache( on_left, batch, &mut *hashmap, offset, - random_state, &mut hashes_buffer, 0, true, null_equality, + &mut hasher, )?; } diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 38da5166e1efe..0a8c3dc79230d 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -27,7 +27,6 @@ use crate::joins::utils::{ use crate::stream::EmptyRecordBatchStream; use crate::{ RecordBatchStream, SendableRecordBatchStream, handle_state, - hash_utils::create_hashes, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, StatefulStreamResult, adjust_indices_by_join_type, apply_join_filter_to_indices, @@ -45,7 +44,7 @@ use datafusion_common::{ }; use datafusion_physical_expr::PhysicalExprRef; -use datafusion_common::hash_utils::RandomState; +use datafusion_common::hash_utils::{CachingArrayHasher, RandomState}; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::{Stream, StreamExt, ready}; @@ -385,6 +384,8 @@ pub(super) struct HashJoinStream { batch_size: usize, /// Scratch space for computing hashes hashes_buffer: Vec, + /// Caches dictionary value hashes by probe key. + hasher: CachingArrayHasher, /// Scratch space for probe indices during hash lookup probe_indices_buffer: Vec, /// Scratch space for build indices during hash lookup @@ -558,12 +559,17 @@ impl HashJoinStream { build_accumulator: Option>, mode: PartitionMode, fetch: Option, - ) -> Self { + ) -> Result { // Create output buffer with coalescing and optional fetch limit. let output_buffer = LimitedBatchCoalescer::new(Arc::clone(&schema), batch_size, fetch); + let fields = on_right + .iter() + .map(|expr| expr.return_field(right.schema().as_ref())) + .collect::>>()?; + let hasher = CachingArrayHasher::new(&Schema::new(fields), random_state.clone()); - Self { + Ok(Self { partition, schema, on_right, @@ -578,6 +584,7 @@ impl HashJoinStream { build_side, batch_size, hashes_buffer, + hasher, probe_indices_buffer: Vec::with_capacity(batch_size), build_indices_buffer: Vec::with_capacity(batch_size), not_in: NotInProbe::default(), @@ -585,7 +592,7 @@ impl HashJoinStream { build_report: BuildReportHandle::new(partition, mode, build_accumulator), mode, output_buffer, - } + }) } /// Returns the next state after the build side has been fully collected @@ -782,11 +789,8 @@ impl HashJoinStream { { self.hashes_buffer.clear(); self.hashes_buffer.resize(batch.num_rows(), 0); - create_hashes( - &keys_values, - &self.random_state, - &mut self.hashes_buffer, - )?; + self.hasher + .create_hashes(&keys_values, &mut self.hashes_buffer)?; matchable_join_keys(&keys_values, self.null_equality) } else { None diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 55fbc2cd11d80..ea6a516773776 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -56,8 +56,9 @@ use arrow_ord::ord::{DynComparator, make_comparator}; use arrow_schema::{DataType, SortOptions, TimeUnit}; use datafusion_common::cast::as_boolean_array; use datafusion_common::float_semantics::normalize_floats; -use datafusion_common::hash_utils::RandomState; +#[cfg(test)] use datafusion_common::hash_utils::create_hashes; +use datafusion_common::hash_utils::{CachingArrayHasher, RandomState}; use datafusion_common::stats::Precision; use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::{ @@ -2349,12 +2350,42 @@ pub fn update_hash( deleted_offset: usize, fifo_hashmap: bool, null_equality: NullEquality, +) -> Result<()> { + let fields = on + .iter() + .map(|expr| expr.return_field(batch.schema_ref())) + .collect::>>()?; + let mut hasher = CachingArrayHasher::new(&Schema::new(fields), random_state.clone()); + update_hash_with_cache( + on, + batch, + hash_map, + offset, + hashes_buffer, + deleted_offset, + fifo_hashmap, + null_equality, + &mut hasher, + ) +} + +#[expect(clippy::too_many_arguments)] +pub(crate) fn update_hash_with_cache( + on: &[PhysicalExprRef], + batch: &RecordBatch, + hash_map: &mut dyn JoinHashMapType, + offset: usize, + hashes_buffer: &mut [u64], + deleted_offset: usize, + fifo_hashmap: bool, + null_equality: NullEquality, + hasher: &mut CachingArrayHasher, ) -> Result<()> { // evaluate the keys let keys_values = evaluate_expressions_to_arrays(on, batch)?; // calculate the hash values - let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?; + let hash_values = hasher.create_hashes(&keys_values, hashes_buffer)?; // For usual JoinHashmap, the implementation is void. hash_map.extend_zero(batch.num_rows());