diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index c9cc48a11dd21..08e85601fed0a 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,73 @@ 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)?; + + 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); + + 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 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] 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 f229cfd513620..ba340b4304bb7 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 @@ -26,15 +26,14 @@ use crate::aggregates::group_values::multi_group_by::{ primitive::PrimitiveGroupValueBuilder, row_backed::RowsGroupColumn, }; use arrow::array::{ - Array, ArrayRef, BooleanBufferBuilder, downcast_primitive, new_empty_array, + Array, ArrayRef, AsArray, 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, create_hashes}; use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::{EmitTo, GroupSelection}; @@ -82,6 +81,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 +218,8 @@ pub struct GroupValuesColumn { /// reused buffer to store hashes hashes_buffer: Vec, - /// Random state for creating hashes - random_state: RandomState, + /// One caching hasher per group column. + hashers: Vec, } /// Buffers to store intermediate results in `vectorized_append` @@ -279,6 +281,13 @@ impl VectorizedOperationBuffers { } } +fn should_cache_dictionary_hashes(array: &dyn Array) -> bool { + // Small dictionaries are cheaper on the existing bulk hashing path. + array + .as_any_dictionary_opt() + .is_some_and(|dictionary| dictionary.values().len() > array.len()) +} + impl GroupValuesColumn { // ======================================================================== // Initialization functions @@ -288,6 +297,7 @@ impl GroupValuesColumn { pub fn try_new(schema: SchemaRef) -> Result { let map = HashTable::with_capacity(0); let group_values = Self::build_group_columns(&schema)?; + let hashers = Self::build_hashers(schema.fields().len()); Ok(Self { schema, map, @@ -297,10 +307,40 @@ impl GroupValuesColumn { map_size: 0, group_values, hashes_buffer: Default::default(), - random_state: crate::aggregates::AGGREGATION_HASH_SEED, + hashers, }) } + fn build_hashers(num_columns: usize) -> Vec { + (0..num_columns) + .map(|_| CachingArrayHasher::new(crate::aggregates::AGGREGATION_HASH_SEED)) + .collect() + } + + fn create_hashes_with_cache( + hashers: &mut [CachingArrayHasher], + group_values: &mut [Box], + cols: &[ArrayRef], + hashes: &mut [u64], + ) -> Result<()> { + if !cols + .iter() + .any(|array| should_cache_dictionary_hashes(array.as_ref())) + { + create_hashes(cols, &crate::aggregates::AGGREGATION_HASH_SEED, hashes)?; + return Ok(()); + } + for (index, ((hasher, array), group_value)) in + hashers.iter_mut().zip(cols).zip(group_values).enumerate() + { + hasher.create_hashes(array.as_ref(), hashes, index != 0)?; + if let Some((values, value_hashes)) = hasher.dictionary_hashes() { + 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 +413,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.hashers, + &mut self.group_values, + cols, + batch_hashes, + )?; for (row, &target_hash) in batch_hashes.iter().enumerate() { let entry = self @@ -474,7 +519,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.hashers, + &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 +1140,12 @@ impl GroupValues for GroupValuesColumn { + self.emit_group_index_list_buffer.allocated_size() + self.vectorized_operation_buffers.size() + self.hashes_buffer.allocated_size() + + self.hashers.allocated_size() + + self + .hashers + .iter() + .map(CachingArrayHasher::size) + .sum::() } fn is_empty(&self) -> bool { @@ -1234,6 +1290,7 @@ 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.hashers = Self::build_hashers(self.schema.fields().len()); // Such structures are only used in `non-streaming` case if !STREAMING { @@ -2070,7 +2127,9 @@ mod tests { use arrow::datatypes::Int8Type; let n_groups = 129_usize; - let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); + let dict_vocab: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..n_groups).map(|i| ["cat", "dog"][i.min(1)]), + )); // Each row has a unique label (forcing a new group) and alternates // between the two dictionary values. Int8 keys are used; only 2 @@ -2093,15 +2152,43 @@ 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 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.hashers[1].dictionary_hashes().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.hashers[1].dictionary_hashes().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)