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] 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..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, } } @@ -182,16 +186,53 @@ 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 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); + } 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 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, @@ -331,13 +372,54 @@ 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() } + fn set_dictionary_hashes(&mut self, values: &ArrayRef, hashes: &[u64]) { + 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], @@ -452,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(); @@ -472,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); } } @@ -598,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) { @@ -610,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 bf784ea5d1591..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 @@ -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, create_hashes}; use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::{EmitTo, GroupSelection}; @@ -82,6 +79,22 @@ 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]) {} + + /// 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; @@ -216,8 +229,14 @@ 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, + + /// 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, } /// Buffers to store intermediate results in `vectorized_append` @@ -288,6 +307,26 @@ 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, + ); + let map_dictionary_groups = !STREAMING + && schema.fields().len() == 1 + && matches!(schema.field(0).data_type(), DataType::Dictionary(_, _)); Ok(Self { schema, map, @@ -297,10 +336,34 @@ impl GroupValuesColumn { map_size: 0, group_values, hashes_buffer: Default::default(), - random_state: crate::aggregates::AGGREGATION_HASH_SEED, + hasher, + dictionary_indices, + map_dictionary_groups, }) } + 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 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); + } + } + Ok(()) + } + /// Build one fresh [`GroupColumn`] per field in the schema. /// /// Used at construction time (`try_new`) and to repopulate the column @@ -313,6 +376,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) } @@ -373,7 +439,13 @@ 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, + &self.dictionary_indices, + cols, + batch_hashes, + )?; for (row, &target_hash) in batch_hashes.iter().enumerate() { let entry = self @@ -474,7 +546,13 @@ 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, + &self.dictionary_indices, + 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`, @@ -1063,6 +1141,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 { @@ -1090,6 +1173,8 @@ 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() } fn is_empty(&self) -> bool { @@ -1247,7 +1332,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 { self.group_index_lists.clear(); @@ -2106,15 +2194,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) @@ -2124,6 +2242,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(); 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);