perf: skip null keys when scattering nullable dictionary hashes

datafusion/common/src/hash_utils.rs

@@ -563,33 +563,55 @@ fn hash_dictionary_scatter<
     hashes_buffer: &mut [u64],
 ) {
     let dict_values = array.values();
-    if HAS_NULL_KEYS {
-        for (hash, key) in hashes_buffer.iter_mut().zip(array.keys().iter()) {
-            if let Some(key) = key {
-                let idx = key.as_usize();
-                if !HAS_NULL_VALUES || 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 {
-                        combine_hashes(dict_hash, *hash)
-                    } else {
-                        dict_hash
-                    };
+    let scatter_row = |hash: &mut u64, key: &K::Native| {
+        let idx = key.as_usize();
+        if !HAS_NULL_VALUES || 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 {
+                combine_hashes(dict_hash, *hash)
+            } else {
+                dict_hash
+            };
+        }
+    };
+    let keys = array.keys().values();
+    if HAS_NULL_KEYS && let Some(key_nulls) = array.keys().nulls() {
+        // Like the `zip` below, stop at the end of `hashes_buffer`. Null slots keep
+        // their hash, and their (possibly out-of-range) keys are never read.
+        let len = hashes_buffer.len().min(keys.len());
+        let (hashes, keys) = (&mut hashes_buffer[..len], &keys[..len]);
+        // Walk the validity bitmap 64 rows at a time: a word with no nulls takes
+        // the same dense loop as the no-null path, which is much cheaper than
+        // visiting each set bit when nulls are sparse
+        let chunks = key_nulls.inner().bit_chunks();
+        let words = chunks
+            .iter()
+            .chain(std::iter::once(chunks.remainder_bits()));
+        for (word_idx, mut word) in words.enumerate() {
+            let start = word_idx * 64;
+            if start >= len {
+                break;
+            }
+            if word == u64::MAX {
+                let end = (start + 64).min(len);
+                for (hash, key) in hashes[start..end].iter_mut().zip(&keys[start..end]) {
+                    scatter_row(hash, key);
+                }
+            } else {
+                while word != 0 {
+                    let i = start + word.trailing_zeros() as usize;
+                    if i >= len {
+                        break;
+                    }
+                    scatter_row(&mut hashes[i], &keys[i]);
+                    word &= word - 1;
                 }
             }
         }
     } 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) {
-                // 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 {
-                    combine_hashes(dict_hash, *hash)
-                } else {
-                    dict_hash
-                };
-            }
+        for (hash, key) in hashes_buffer.iter_mut().zip(keys) {
+            scatter_row(hash, key);
         }
     }
 }
@@ -1844,6 +1866,179 @@ mod tests {
         assert_ne!(dict_hashes[0], dict_hashes[2]);
     }
 
+    #[test]
+    // Tests actual values of hashes, which are different if forcing collisions
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn create_hashes_for_dict_arrays_with_null_keys() {
+        check_dict_null_keys::<Int8Type>();
+        check_dict_null_keys::<Int32Type>();
+        check_dict_null_keys::<Int64Type>();
+    }
+
+    /// Compares hashes of sliced dictionaries with null keys against the
+    /// equivalent `StringArray`, for one and two columns
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn check_dict_null_keys<K: ArrowDictionaryKeyType>() {
+        let random_state = RandomState::with_seed(0);
+        let value_sets: [&[Option<&str>]; 2] = [
+            &[Some("alpha"), Some("beta"), Some("gamma")],
+            &[Some("alpha"), None, Some("gamma")],
+        ];
+        // Scattered nulls, then an all-valid word, an all-null word, alternating
+        let is_null = |row: usize| match row / 64 {
+            0 => row % 5 == 1,
+            1 => false,
+            2 => true,
+            _ => row.is_multiple_of(2),
+        };
+        for values in value_sets {
+            // Out-of-range keys hidden behind null slots must never be read
+            let invalid_keys = [usize::MAX, values.len(), 127];
+            for offset in [0, 3, 5, 64] {
+                for len in [0, 1, 7, 8, 9, 63, 64, 65, 127, 128, 129] {
+                    let rows = 0..offset + len;
+                    let keys = rows
+                        .clone()
+                        .map(|row| {
+                            let key = if is_null(row) {
+                                invalid_keys[row % invalid_keys.len()]
+                            } else {
+                                row % values.len()
+                            };
+                            K::Native::usize_as(key)
+                        })
+                        .collect::<Vec<_>>();
+                    let validity =
+                        rows.clone().map(|row| !is_null(row)).collect::<Vec<_>>();
+                    let dict = DictionaryArray::<K>::try_new(
+                        PrimitiveArray::new(
+                            ScalarBuffer::from(keys),
+                            Some(NullBuffer::from(validity)),
+                        ),
+                        Arc::new(values.iter().copied().collect::<StringArray>()),
+                    )
+                    .unwrap()
+                    .slice(offset, len);
+                    let strings = rows
+                        .skip(offset)
+                        .map(|row| {
+                            if is_null(row) {
+                                None
+                            } else {
+                                values[row % values.len()]
+                            }
+                        })
+                        .collect::<StringArray>();
+                    let prefix = Int32Array::from_iter_values(0..len as i32);
+
+                    let mut expected = vec![0; len];
+                    create_hashes([&strings as &dyn Array], &random_state, &mut expected)
+                        .unwrap();
+                    let mut actual = vec![0; len];
+                    create_hashes([&dict as &dyn Array], &random_state, &mut actual)
+                        .unwrap();
+                    assert_eq!(actual, expected, "single, offset {offset}, len {len}");
+
+                    // A shorter buffer hashes only its prefix
+                    let mut short = vec![0; len / 2];
+                    create_hashes([&dict as &dyn Array], &random_state, &mut short)
+                        .unwrap();
+                    assert_eq!(
+                        short,
+                        expected[..len / 2],
+                        "short, offset {offset}, len {len}"
+                    );
+
+                    let mut expected = vec![0; len];
+                    create_hashes(
+                        [&prefix as &dyn Array, &strings],
+                        &random_state,
+                        &mut expected,
+                    )
+                    .unwrap();
+                    let mut actual = vec![0; len];
+                    create_hashes(
+                        [&prefix as &dyn Array, &dict],
+                        &random_state,
+                        &mut actual,
+                    )
+                    .unwrap();
+                    assert_eq!(actual, expected, "multiple, offset {offset}, len {len}");
+                }
+            }
+        }
+
+        // An empty dictionary is valid when every key is null
+        let dict = DictionaryArray::<K>::try_new(
+            PrimitiveArray::new(
+                ScalarBuffer::from(vec![K::Native::usize_as(usize::MAX); 64]),
+                Some(NullBuffer::new_null(64)),
+            ),
+            Arc::new(StringArray::new_null(0)),
+        )
+        .unwrap();
+        let mut hashes = vec![0; 64];
+        create_hashes([&dict as &dyn Array], &random_state, &mut hashes).unwrap();
+        assert_eq!(hashes, vec![0; 64]);
+        let prefix = Int32Array::from_iter_values(0..64);
+        let mut expected = vec![0; 64];
+        create_hashes([&prefix as &dyn Array], &random_state, &mut expected).unwrap();
+        let mut actual = vec![0; 64];
+        create_hashes([&prefix as &dyn Array, &dict], &random_state, &mut actual)
+            .unwrap();
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    // Tests actual values of hashes, which are different if forcing collisions
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn create_hashes_for_dict_short_buffer_before_all_valid_word() {
+        // The validity word at row 64 is all valid and starts past the end of
+        // the shorter buffers, and the last row is null
+        let rows = 200;
+        let keys = Int32Array::from_iter((0..rows).map(|row| (row != 199).then_some(0)));
+        let dict = DictionaryArray::try_new(
+            keys,
+            Arc::new(StringArray::from(vec!["alpha"])) as ArrayRef,
+        )
+        .unwrap();
+        let strings = (0..rows)
+            .map(|row| (row != 199).then_some("alpha"))
+            .collect::<StringArray>();
+        let prefix = Int32Array::from_iter_values(0..rows);
+        let random_state = RandomState::with_seed(0);
+
+        // The equivalent `StringArray` only hashes into a full-length buffer
+        let mut expected_single = vec![0; rows as usize];
+        create_hashes(
+            [&strings as &dyn Array],
+            &random_state,
+            &mut expected_single,
+        )
+        .unwrap();
+        let mut expected_multiple = vec![0; rows as usize];
+        create_hashes(
+            [&prefix as &dyn Array, &strings],
+            &random_state,
+            &mut expected_multiple,
+        )
+        .unwrap();
+
+        // A shorter buffer hashes only its prefix. Only a single column can have
+        // one: any other column must be as long as the buffer
+        for len in [0, 1, 10, 63, 64, 65, 100, 128, 129, 199, 200] {
+            let mut actual = vec![0; len];
+            create_hashes([&dict as &dyn Array], &random_state, &mut actual).unwrap();
+            assert_eq!(actual, expected_single[..len], "buffer len {len}");
+        }
+
+        // The same dictionary as a later column of a full-length buffer
+        let mut actual = vec![0; rows as usize];
+        create_hashes([&prefix as &dyn Array, &dict], &random_state, &mut actual)
+            .unwrap();
+        assert_eq!(actual, expected_multiple);
+    }
+
     #[test]
     // Tests actual values of hashes, which are different if forcing collisions
     #[cfg(not(feature = "force_hash_collisions"))]