diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index c9cc48a11dd21..1bc1fb360f57a 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -96,6 +96,20 @@ pub fn combine_hashes(l: u64, r: u64) -> u64 { /// to this size. const MAX_BUFFER_SIZE: usize = 524_288; +#[cfg(not(feature = "force_hash_collisions"))] +fn dictionary_nulls( + array: &DictionaryArray, +) -> (bool, bool) { + ( + array.keys().null_count() != 0, + array.values().null_count() != 0, + ) +} + +fn dictionary_overwrites_hash_buffer(array: &dyn Array) -> bool { + matches!(array.data_type(), DataType::Dictionary(_, _)) && !array.is_nullable() +} + thread_local! { /// Thread-local buffer for hash computations to avoid repeated allocations. /// The buffer is reused across calls and truncated if it exceeds MAX_BUFFER_SIZE. @@ -158,8 +172,14 @@ where let mut buffer = cell.try_borrow_mut() .map_err(|_| _internal_datafusion_err!("with_hashes cannot be called reentrantly on the same thread"))?; - // Ensure buffer has sufficient length, clearing old values - buffer.clear(); + // A dictionary with non-null keys and values overwrites every output + // slot, so retain the initialized storage instead of zeroing it first. + let reuse_buffer = iter + .peek() + .is_some_and(|array| dictionary_overwrites_hash_buffer(array.as_dyn_array())); + if !reuse_buffer { + buffer.clear(); + } buffer.resize(required_size, 0); // Create hashes in the buffer - this consumes the iterator @@ -544,6 +564,65 @@ fn hash_generic_byte_view_array( } } +/// Scatters non-null dictionary hashes in independent lanes to overlap key-indexed loads. +#[cfg(not(feature = "force_hash_collisions"))] +#[inline(always)] +fn scatter_non_null_dictionary( + array: &DictionaryArray, + dict_hashes: &[u64], + hashes: &mut [u64], +) { + const LANES: usize = 8; + debug_assert_eq!(dict_hashes.len(), array.values().len()); + debug_assert_eq!(dictionary_nulls(array), (false, false)); + let keys = array.keys().values(); + let len = keys.len().min(hashes.len()); + debug_assert_eq!( + len, + hashes.len(), + "dictionary scatter must overwrite every hash slot" + ); + // `values()` accounts for key-array slices, and all key widths share the + // same `as_usize` lookup below. + let keys = &keys[..len]; + let hashes = &mut hashes[..len]; + + let (key_chunks, key_tail) = keys.as_chunks::(); + let (hash_chunks, hash_tail) = hashes.as_chunks_mut::(); + for (keys, hashes) in key_chunks.iter().zip(hash_chunks) { + let gathered: [u64; LANES] = std::array::from_fn(|lane| unsafe { + let key = keys[lane]; + // SAFETY: Dictionary keys are in bounds for values, which has the + // same length as dict_hashes. + *dict_hashes.get_unchecked(key.as_usize()) + }); + if MULTI_COL { + for (hash, dict_hash) in hashes.iter_mut().zip(gathered) { + update_dictionary_hash::(hash, dict_hash); + } + } else { + *hashes = gathered; + } + } + + for (hash, key) in hash_tail.iter_mut().zip(key_tail) { + // SAFETY: Dictionary keys are in bounds for values, which has the + // same length as dict_hashes. + let dict_hash = unsafe { *dict_hashes.get_unchecked(key.as_usize()) }; + update_dictionary_hash::(hash, dict_hash); + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +#[inline(always)] +fn update_dictionary_hash(hash: &mut u64, dict_hash: u64) { + *hash = if MULTI_COL { + combine_hashes(dict_hash, *hash) + } else { + dict_hash + }; +} + /// Scatter precomputed dictionary value hashes to key positions. /// /// Uses const generics to eliminate runtime branching in the hot loop: @@ -562,6 +641,10 @@ fn hash_dictionary_scatter< dict_hashes: &[u64], hashes_buffer: &mut [u64], ) { + if !HAS_NULL_KEYS && !HAS_NULL_VALUES { + scatter_non_null_dictionary::(array, dict_hashes, hashes_buffer); + return; + } let dict_values = array.values(); if HAS_NULL_KEYS { for (hash, key) in hashes_buffer.iter_mut().zip(array.keys().iter()) { @@ -581,7 +664,7 @@ fn hash_dictionary_scatter< } else { for (hash, key) in hashes_buffer.iter_mut().zip(array.keys().values()) { let idx = key.as_usize(); - if !HAS_NULL_VALUES || dict_values.is_valid(idx) { + if dict_values.is_valid(idx) { // SAFETY: `idx` is a valid Arrow dict key, in-bounds for `dict_hashes`. let dict_hash = unsafe { *dict_hashes.get_unchecked(idx) }; *hash = if MULTI_COL { @@ -601,8 +684,7 @@ fn dispatch_dictionary_scatter( hashes_buffer: &mut [u64], multi_col: bool, ) { - let has_null_keys = array.keys().null_count() != 0; - let has_null_values = array.values().null_count() != 0; + let (has_null_keys, has_null_values) = dictionary_nulls(array); match (has_null_keys, has_null_values, multi_col) { (false, false, false) => hash_dictionary_scatter::( @@ -2284,6 +2366,78 @@ mod tests { assert_eq!(result, expected_hashes); } + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_non_null_dictionary_hash_fast_path() { + let random_state = RandomState::with_seed(0); + let keys = vec![2, 0, 1, 2, 1, 0, 2, 0, 1, 2, 0]; + let non_null: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int32Array::from(keys.clone()), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ) + .unwrap(), + ); + let null_values: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])), + ) + .unwrap(), + ); + let null_keys: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![ + Some(2), + None, + Some(1), + Some(2), + None, + Some(0), + Some(2), + None, + Some(1), + Some(2), + Some(0), + ]), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ) + .unwrap(), + ); + + let assert_matches = |dictionary: &ArrayRef| { + let strings = arrow::compute::cast(dictionary, &DataType::Utf8).unwrap(); + let mut expected = vec![0; dictionary.len()]; + create_hashes([strings], &random_state, &mut expected).unwrap(); + with_hashes([dictionary], &random_state, |hashes| { + assert_eq!(hashes, expected); + Ok(()) + }) + .unwrap(); + }; + + let strings = arrow::compute::cast(&non_null, &DataType::Utf8).unwrap(); + let mut expected = vec![0; strings.len()]; + create_hashes([&strings], &random_state, &mut expected).unwrap(); + let mut prefix = vec![0; non_null.len() - 1]; + create_hashes([&non_null], &random_state, &mut prefix).unwrap(); + assert_eq!(prefix, expected[..prefix.len()]); + + let seed: ArrayRef = + Arc::new(Int32Array::from_iter_values(0..non_null.len() as i32)); + let mut expected = vec![0; seed.len()]; + create_hashes([&seed, &strings], &random_state, &mut expected).unwrap(); + with_hashes([&seed, &non_null], &random_state, |hashes| { + assert_eq!(hashes, expected); + Ok(()) + }) + .unwrap(); + + assert_matches(&non_null); + assert_matches(&null_values); + assert_matches(&null_keys); + } + #[test] fn test_with_hashes_multi_column() { let int_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));