From ab7a91808e6db82783661fca3e1df2297947ea06 Mon Sep 17 00:00:00 2001 From: Lakshya Sharma <229351168+lakshaydevoo-arch@users.noreply.github.com> Date: Sat, 31 Oct 2026 09:29:55 +0000 Subject: [PATCH 1/2] perf: skip null keys when scattering nullable dictionary hashes --- datafusion/common/benches/with_hashes.rs | 60 ++++++++- datafusion/common/src/hash_utils.rs | 153 +++++++++++++++++++++-- 2 files changed, 194 insertions(+), 19 deletions(-) diff --git a/datafusion/common/benches/with_hashes.rs b/datafusion/common/benches/with_hashes.rs index 7fc4a29fc55ec..3a22325e58781 100644 --- a/datafusion/common/benches/with_hashes.rs +++ b/datafusion/common/benches/with_hashes.rs @@ -178,8 +178,13 @@ fn create_null_mask(len: usize) -> NullBuffer where StandardUniform: Distribution, { + create_null_mask_with_density(len, 0.03).expect("should be nulls in buffer") +} + +/// Returns a random null mask where each slot is null with probability +/// `null_density`, or `None` if no slot ended up null +fn create_null_mask_with_density(len: usize, null_density: f32) -> Option { let mut rng = make_rng(); - let null_density = 0.03; let mut builder = NullBufferBuilder::new(len); for _ in 0..len { if rng.random::() < null_density { @@ -188,7 +193,21 @@ where builder.append_non_null(); } } - builder.finish().expect("should be nulls in buffer") + builder.finish() +} + +/// Returns `dict` with a new random null mask on its keys, where each key is null +/// with probability `null_density`. Any existing key nulls are replaced, and the +/// key and value buffers are shared rather than copied +fn with_null_keys( + dict: &DictionaryArray, + null_density: f32, +) -> ArrayRef { + let keys = PrimitiveArray::::new( + dict.keys().values().clone(), + create_null_mask_with_density(dict.len(), null_density), + ); + Arc::new(DictionaryArray::new(keys, Arc::clone(dict.values()))) } // Returns a new array that is the same as array, but with nulls @@ -290,9 +309,15 @@ impl StringPool { &self, array_length: usize, ) -> ArrayRef { - Arc::new(DictionaryArray::::from_iter( - self.iter_strings(array_length), - )) + Arc::new(self.typed_dictionary_array::(array_length)) + } + + /// Like [`Self::dictionary_array`], but returns the concrete `DictionaryArray` + fn typed_dictionary_array( + &self, + array_length: usize, + ) -> DictionaryArray { + DictionaryArray::::from_iter(self.iter_strings(array_length)) } } @@ -645,10 +670,33 @@ fn mixed_columns_benchmark(c: &mut Criterion) { }); } +/// Int32 dictionaries with 50% and 97% null keys. 0% and 3% are covered by +/// the existing `dictionary_utf8_int32` "no nulls" and "nulls" cases in +/// `criterion_benchmark` +fn dictionary_benchmark(c: &mut Criterion) { + let pool = StringPool::new(100, 64); + let dict = pool.typed_dictionary_array::(BATCH_SIZE); + let name = "dictionary_utf8_int32"; + for null_percent in [50, 97] { + let array = with_null_keys(&dict, null_percent as f32 / 100.0); + c.bench_function(&format!("{name}: single, {null_percent}% null keys"), |b| { + do_hash_test(b, std::slice::from_ref(&array)); + }); + c.bench_function( + &format!("{name}: multiple, {null_percent}% null keys"), + |b| { + let arrays = vec![array.clone(), array.clone(), array.clone()]; + do_hash_test(b, &arrays); + }, + ); + } +} + criterion_group!( benches, criterion_benchmark, sliced_array_benchmark, - mixed_columns_benchmark + mixed_columns_benchmark, + dictionary_benchmark ); criterion_main!(benches); 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"))] From 26d37c7dd18f59237d820371b1892b6652aaacdd Mon Sep 17 00:00:00 2001 From: Lakshya Sharma <229351168+lakshaydevoo-arch@users.noreply.github.com> Date: Mon, 2 Nov 2026 07:28:45 +0000 Subject: [PATCH 2/2] refactor: separate dictionary benchmarks into dedicated PR --- datafusion/common/benches/with_hashes.rs | 60 +++--------------------- 1 file changed, 6 insertions(+), 54 deletions(-) diff --git a/datafusion/common/benches/with_hashes.rs b/datafusion/common/benches/with_hashes.rs index 3a22325e58781..7fc4a29fc55ec 100644 --- a/datafusion/common/benches/with_hashes.rs +++ b/datafusion/common/benches/with_hashes.rs @@ -178,13 +178,8 @@ fn create_null_mask(len: usize) -> NullBuffer where StandardUniform: Distribution, { - create_null_mask_with_density(len, 0.03).expect("should be nulls in buffer") -} - -/// Returns a random null mask where each slot is null with probability -/// `null_density`, or `None` if no slot ended up null -fn create_null_mask_with_density(len: usize, null_density: f32) -> Option { let mut rng = make_rng(); + let null_density = 0.03; let mut builder = NullBufferBuilder::new(len); for _ in 0..len { if rng.random::() < null_density { @@ -193,21 +188,7 @@ fn create_null_mask_with_density(len: usize, null_density: f32) -> Option( - dict: &DictionaryArray, - null_density: f32, -) -> ArrayRef { - let keys = PrimitiveArray::::new( - dict.keys().values().clone(), - create_null_mask_with_density(dict.len(), null_density), - ); - Arc::new(DictionaryArray::new(keys, Arc::clone(dict.values()))) + builder.finish().expect("should be nulls in buffer") } // Returns a new array that is the same as array, but with nulls @@ -309,15 +290,9 @@ impl StringPool { &self, array_length: usize, ) -> ArrayRef { - Arc::new(self.typed_dictionary_array::(array_length)) - } - - /// Like [`Self::dictionary_array`], but returns the concrete `DictionaryArray` - fn typed_dictionary_array( - &self, - array_length: usize, - ) -> DictionaryArray { - DictionaryArray::::from_iter(self.iter_strings(array_length)) + Arc::new(DictionaryArray::::from_iter( + self.iter_strings(array_length), + )) } } @@ -670,33 +645,10 @@ fn mixed_columns_benchmark(c: &mut Criterion) { }); } -/// Int32 dictionaries with 50% and 97% null keys. 0% and 3% are covered by -/// the existing `dictionary_utf8_int32` "no nulls" and "nulls" cases in -/// `criterion_benchmark` -fn dictionary_benchmark(c: &mut Criterion) { - let pool = StringPool::new(100, 64); - let dict = pool.typed_dictionary_array::(BATCH_SIZE); - let name = "dictionary_utf8_int32"; - for null_percent in [50, 97] { - let array = with_null_keys(&dict, null_percent as f32 / 100.0); - c.bench_function(&format!("{name}: single, {null_percent}% null keys"), |b| { - do_hash_test(b, std::slice::from_ref(&array)); - }); - c.bench_function( - &format!("{name}: multiple, {null_percent}% null keys"), - |b| { - let arrays = vec![array.clone(), array.clone(), array.clone()]; - do_hash_test(b, &arrays); - }, - ); - } -} - criterion_group!( benches, criterion_benchmark, sliced_array_benchmark, - mixed_columns_benchmark, - dictionary_benchmark + mixed_columns_benchmark ); criterion_main!(benches);