Skip to content

Commit 1014e2e

Browse files
committed
perf: hash sparse dictionary slices by logical rows
1 parent 6213ad2 commit 1014e2e

2 files changed

Lines changed: 91 additions & 1 deletion

File tree

‎datafusion/common/benches/with_hashes.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,31 @@ fn sliced_array_benchmark(c: &mut Criterion) {
409409
}
410410
}
411411

412+
/// Benchmark dictionary arrays whose keys reference a much larger values array.
413+
///
414+
/// This shape is produced when a large aggregate dictionary is split into
415+
/// smaller record batches: slicing changes the keys but retains the complete
416+
/// dictionary values array.
417+
fn sparse_dictionary_benchmark(c: &mut Criterion) {
418+
for ratio in [1, 2, 4, 8, 32, 128] {
419+
let values_len = BATCH_SIZE * ratio;
420+
let values: ArrayRef = Arc::new(GenericStringArray::<i32>::from_iter_values(
421+
(0..values_len).map(|i| format!("value-{i:016x}-with-a-long-suffix")),
422+
));
423+
let keys =
424+
Int32Array::from_iter_values((0..BATCH_SIZE).map(|i| (i * ratio) as i32));
425+
let dictionary: ArrayRef =
426+
Arc::new(DictionaryArray::<Int32Type>::new(keys, values));
427+
428+
c.bench_function(
429+
&format!(
430+
"dictionary_utf8_sparse_backing: {BATCH_SIZE} rows / {values_len} values"
431+
),
432+
|b| do_hash_test(b, std::slice::from_ref(&dictionary)),
433+
);
434+
}
435+
}
436+
412437
fn do_hash_test_with_len(b: &mut Bencher, arrays: &[ArrayRef], expected_len: usize) {
413438
let state = RandomState::default();
414439
b.iter(|| {
@@ -662,6 +687,7 @@ criterion_group!(
662687
benches,
663688
criterion_benchmark,
664689
sliced_array_benchmark,
690+
sparse_dictionary_benchmark,
665691
mixed_columns_benchmark
666692
);
667693
criterion_main!(benches);

‎datafusion/common/src/hash_utils.rs‎

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -631,10 +631,26 @@ fn hash_dictionary<K: ArrowDictionaryKeyType>(
631631
hashes_buffer: &mut [u64],
632632
multi_col: bool,
633633
) -> Result<()> {
634+
let dict_values = array.values();
635+
636+
// Slicing a dictionary only slices its keys and retains the complete values
637+
// array. Hashing every retained value can therefore be much more expensive
638+
// than hashing the logical rows in a small slice. Expand only for this
639+
// sparse-backing case; normal dictionaries keep the hash-once-and-scatter
640+
// path below.
641+
if dict_values.len() > array.len().saturating_mul(4) {
642+
let logical_values = take(dict_values.as_ref(), array.keys(), None)?;
643+
return hash_single_array(
644+
logical_values.as_ref(),
645+
random_state,
646+
hashes_buffer,
647+
multi_col,
648+
);
649+
}
650+
634651
// Hash each dictionary value once, and then use that computed
635652
// hash for each key value to avoid a potentially expensive
636653
// redundant hashing for large dictionary elements (e.g. strings)
637-
let dict_values = array.values();
638654
let mut dict_hashes = vec![0; dict_values.len()];
639655
create_hashes([dict_values], random_state, &mut dict_hashes)?;
640656
dispatch_dictionary_scatter(array, &dict_hashes, hashes_buffer, multi_col);
@@ -2164,6 +2180,54 @@ mod tests {
21642180
assert_ne!(one_col_hashes, two_col_hashes);
21652181
}
21662182

2183+
#[test]
2184+
#[cfg(not(feature = "force_hash_collisions"))]
2185+
fn sparse_dictionary_hashes_match_expanded_values() {
2186+
let values = (0..64)
2187+
.map(|i| (i != 7).then(|| format!("value-{i}")))
2188+
.collect::<Vec<_>>();
2189+
let values: ArrayRef =
2190+
Arc::new(StringArray::from_iter(values.iter().map(Option::as_deref)));
2191+
let keys = Int32Array::from(vec![Some(63), None, Some(7), Some(2)]);
2192+
let logical_values = take(values.as_ref(), &keys, None).unwrap();
2193+
let dictionary: ArrayRef =
2194+
Arc::new(DictionaryArray::<Int32Type>::new(keys, Arc::clone(&values)));
2195+
let prefix: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
2196+
let random_state = RandomState::with_seed(0);
2197+
2198+
let mut dictionary_hashes = vec![0; dictionary.len()];
2199+
create_hashes(
2200+
&[Arc::clone(&dictionary)],
2201+
&random_state,
2202+
&mut dictionary_hashes,
2203+
)
2204+
.unwrap();
2205+
let mut expanded_hashes = vec![0; logical_values.len()];
2206+
create_hashes(
2207+
&[Arc::clone(&logical_values)],
2208+
&random_state,
2209+
&mut expanded_hashes,
2210+
)
2211+
.unwrap();
2212+
assert_eq!(dictionary_hashes, expanded_hashes);
2213+
2214+
dictionary_hashes.fill(0);
2215+
create_hashes(
2216+
&[Arc::clone(&prefix), dictionary],
2217+
&random_state,
2218+
&mut dictionary_hashes,
2219+
)
2220+
.unwrap();
2221+
expanded_hashes.fill(0);
2222+
create_hashes(
2223+
&[prefix, logical_values],
2224+
&random_state,
2225+
&mut expanded_hashes,
2226+
)
2227+
.unwrap();
2228+
assert_eq!(dictionary_hashes, expanded_hashes);
2229+
}
2230+
21672231
#[test]
21682232
fn test_create_hashes_from_arrays() {
21692233
let int_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));

0 commit comments

Comments
 (0)