perf: speed up non-null dictionary hash scatter

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<K: ArrowDictionaryKeyType>(
+    array: &DictionaryArray<K>,
+) -> (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<T: ByteViewType>(
     }
 }
 
+/// 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<K: ArrowDictionaryKeyType, const MULTI_COL: bool>(
+    array: &DictionaryArray<K>,
+    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::<LANES>();
+    let (hash_chunks, hash_tail) = hashes.as_chunks_mut::<LANES>();
+    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::<MULTI_COL>(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::<MULTI_COL>(hash, dict_hash);
+    }
+}
+
+#[cfg(not(feature = "force_hash_collisions"))]
+#[inline(always)]
+fn update_dictionary_hash<const MULTI_COL: bool>(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::<K, MULTI_COL>(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<K: ArrowDictionaryKeyType>(
     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::<K, false, false, false>(
@@ -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::<Int32Type>::try_new(
+                Int32Array::from(keys.clone()),
+                Arc::new(StringArray::from(vec!["a", "b", "c"])),
+            )
+            .unwrap(),
+        );
+        let null_values: ArrayRef = Arc::new(
+            DictionaryArray::<Int32Type>::try_new(
+                Int32Array::from(keys),
+                Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])),
+            )
+            .unwrap(),
+        );
+        let null_keys: ArrayRef = Arc::new(
+            DictionaryArray::<Int32Type>::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]));