diff --git a/datafusion/physical-plan/benches/dictionary_group_values.rs b/datafusion/physical-plan/benches/dictionary_group_values.rs index ded52aebd1100..c0dafccd7316d 100644 --- a/datafusion/physical-plan/benches/dictionary_group_values.rs +++ b/datafusion/physical-plan/benches/dictionary_group_values.rs @@ -29,7 +29,7 @@ use criterion::{ }; use datafusion_expr::EmitTo; use datafusion_physical_plan::aggregates::group_values::new_group_values; -use datafusion_physical_plan::aggregates::order::GroupOrdering; +use datafusion_physical_plan::aggregates::order::{GroupOrdering, GroupOrderingFull}; use rand::rngs::StdRng; use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; @@ -172,5 +172,85 @@ fn bench_repeated_intern_emit(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_intern_emit, bench_repeated_intern_emit); +// GroupOrdering::Full -> GroupValuesColumn::: scalar append_val/equal_to path. +fn bench_scalar_append_equal(c: &mut Criterion) { + let mut group = c.benchmark_group("dict_scalar_append_equal"); + let schema = dict_schema(); + let null_density = 0.1; + let size = SIZES[1]; + + let mut cards = CARDS_RELATIVE.to_vec(); + cards.push(size); + for cardinality in cards { + let array = make_dict(size, cardinality, null_density, SEED); + group.throughput(Throughput::Elements(size as u64)); + group.bench_function( + bench_id("scalar_append_equal", size, cardinality, null_density), + |b| { + b.iter_batched_ref( + || { + ( + new_group_values( + schema.clone(), + &GroupOrdering::Full(GroupOrderingFull::new()), + ) + .unwrap(), + Vec::::with_capacity(size), + ) + }, + |(gv, groups)| { + gv.intern(std::slice::from_ref(&array), groups).unwrap(); + black_box(&*groups); + black_box(gv.emit(EmitTo::All).unwrap()); + }, + BatchSize::SmallInput, + ); + }, + ); + } + group.finish(); +} + +// EmitTo::First exercises the take-n path; two interns + partial emit per iteration. +fn bench_take_n(c: &mut Criterion) { + let mut group = c.benchmark_group("dict_take_n"); + let schema = dict_schema(); + let null_density = 0.10; + let size = SIZES[1]; + + let mut cards = CARDS_RELATIVE.to_vec(); + cards.push(size); + for cardinality in cards { + let batch = make_dict(size, cardinality, null_density, SEED); + group.throughput(Throughput::Elements((size * N_BATCHES) as u64)); + group.bench_function(bench_id("take_n", size, cardinality, null_density), |b| { + b.iter_batched_ref( + || { + ( + new_group_values(schema.clone(), &GroupOrdering::None).unwrap(), + Vec::::with_capacity(size), + ) + }, + |(gv, groups)| { + for _ in 0..N_BATCHES { + gv.intern(std::slice::from_ref(&batch), groups).unwrap(); + black_box(&*groups); + black_box(gv.emit(EmitTo::First(size / 2)).unwrap()); + } + black_box(gv.emit(EmitTo::First(gv.len())).unwrap()); + }, + BatchSize::SmallInput, + ); + }); + } + group.finish(); +} + +criterion_group!( + benches, + bench_intern_emit, + bench_repeated_intern_emit, + bench_scalar_append_equal, + bench_take_n +); criterion_main!(benches); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index 501b13d0cd183..d0c167e1d2d59 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -57,6 +57,9 @@ pub struct DictionaryGroupValuesColumn val_to_inner: Vec, /// Reusable hash buffer for the dictionary values array. val_hashes: Vec, + /// The last `dict.values()` Arc hashed in `append_val`. When the incoming + /// values array is `ptr_eq` to this, `val_hashes` can be reused directly. + cached_values: Option, _phantom: PhantomData, } @@ -73,6 +76,7 @@ impl DictionaryGroupValuesColumn { random_state: AGGREGATION_HASH_SEED, val_to_inner: Vec::default(), val_hashes: Vec::default(), + cached_values: None, _phantom: PhantomData, } } @@ -159,6 +163,7 @@ impl DictionaryGroupValuesColumn { } fn hash_values(&mut self, values: &ArrayRef) { + self.cached_values = None; self.val_hashes.clear(); self.val_hashes.resize(values.len(), 0); create_hashes( @@ -296,9 +301,25 @@ impl GroupColumn } Some(val_idx) => { let dict_values = dict.values(); - let single = dict_values.slice(val_idx, 1); - self.hash_values(&single); - self.find_or_insert_value(dict_values, val_idx, self.val_hashes[0])? + // check if the dictionary values array we are hashing was already seen. + // if its arc was already stored we dont need to rehash the entire array again + // if its new hash the entire array and store an arc ptr for future use + let cache_hit = self + .cached_values + .as_ref() + .is_some_and(|c| Arc::ptr_eq(c, dict_values)); + if !cache_hit { + self.val_hashes.clear(); + self.val_hashes.resize(dict_values.len(), 0); + create_hashes( + std::slice::from_ref(dict_values), + &self.random_state, + &mut self.val_hashes, + ) + .unwrap(); + self.cached_values = Some(Arc::clone(dict_values)); + } + self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx])? } }; self.group_to_inner.push(inner_slot);