diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index c9cc48a11dd21..9d5770350caa4 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -563,19 +563,23 @@ 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 - }; - } + if HAS_NULL_KEYS && let Some(key_nulls) = array.keys().nulls() { + // Visit only valid key slots: null slots keep their hash, and their + // (possibly out-of-range) key values are never read + let keys = array.keys().values(); + // Like the `zip` in the loop below, stop at the end of `hashes_buffer` + let len = hashes_buffer.len(); + for i in key_nulls.valid_indices().take_while(|&i| i < len) { + let idx = keys[i].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) }; + let hash = &mut hashes_buffer[i]; + *hash = if MULTI_COL { + combine_hashes(dict_hash, *hash) + } else { + dict_hash + }; } } } else { @@ -1838,6 +1842,129 @@ 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::(); + check_dict_null_keys::(); + check_dict_null_keys::(); + } + + /// 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() { + 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::>(); + let validity = + rows.clone().map(|row| !is_null(row)).collect::>(); + let dict = DictionaryArray::::try_new( + PrimitiveArray::new( + ScalarBuffer::from(keys), + Some(NullBuffer::from(validity)), + ), + Arc::new(values.iter().copied().collect::()), + ) + .unwrap() + .slice(offset, len); + let strings = rows + .skip(offset) + .map(|row| { + if is_null(row) { + None + } else { + values[row % values.len()] + } + }) + .collect::(); + 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::::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"))]