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]