Skip to content

Commit 54ba4a1

Browse files
author
B Vadlamani
committed
fix_pr_comments
1 parent 7d3e01f commit 54ba4a1

2 files changed

Lines changed: 182 additions & 6 deletions

File tree

‎datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,10 @@ mod native;
2323
pub use bytes::BytesDistinctCountAccumulator;
2424
pub use bytes::BytesViewDistinctCountAccumulator;
2525
pub use dict::DictionaryCountAccumulator;
26-
<<<<<<< HEAD
26+
pub use groups::PrimitiveDistinctCountGroupsAccumulator;
2727
pub use native::Bitmap65536DistinctCountAccumulator;
2828
pub use native::Bitmap65536DistinctCountAccumulatorI16;
2929
pub use native::BoolArray256DistinctCountAccumulator;
3030
pub use native::BoolArray256DistinctCountAccumulatorI8;
31-
=======
32-
pub use groups::PrimitiveDistinctCountGroupsAccumulator;
33-
>>>>>>> f982d8de2 (add count distinct group benchmarks)
3431
pub use native::FloatDistinctCountAccumulator;
3532
pub use native::PrimitiveDistinctCountAccumulator;

‎datafusion/functions-aggregate/benches/count_distinct.rs‎

Lines changed: 181 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@
1818
use std::sync::Arc;
1919

2020
use arrow::array::{
21-
Array, ArrayRef, Int8Array, Int16Array, Int64Array, UInt8Array, UInt16Array,
21+
Array, ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array, UInt8Array,
22+
UInt16Array, UInt32Array,
2223
};
2324
use arrow::datatypes::{DataType, Field, Schema};
2425
use criterion::{Criterion, criterion_group, criterion_main};
@@ -87,6 +88,20 @@ fn create_i16_array(n_distinct: usize) -> Int16Array {
8788
.collect()
8889
}
8990

91+
fn create_u32_array(n_distinct: usize) -> UInt32Array {
92+
let mut rng = StdRng::seed_from_u64(42);
93+
(0..BATCH_SIZE)
94+
.map(|_| Some(rng.random_range(0..n_distinct as u32)))
95+
.collect()
96+
}
97+
98+
fn create_i32_array(n_distinct: usize) -> Int32Array {
99+
let mut rng = StdRng::seed_from_u64(42);
100+
(0..BATCH_SIZE)
101+
.map(|_| Some(rng.random_range(0..n_distinct as i32)))
102+
.collect()
103+
}
104+
90105
fn prepare_args(data_type: DataType) -> (Arc<Schema>, AccumulatorArgs<'static>) {
91106
let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)]));
92107
let schema_leaked: &'static Schema = Box::leak(Box::new((*schema).clone()));
@@ -172,6 +187,33 @@ fn count_distinct_benchmark(c: &mut Criterion) {
172187
.unwrap()
173188
})
174189
});
190+
191+
// 32-bit integer types
192+
for pct in [80, 99] {
193+
let n_distinct = BATCH_SIZE * pct / 100;
194+
195+
// UInt32
196+
let values = Arc::new(create_u32_array(n_distinct)) as ArrayRef;
197+
c.bench_function(&format!("count_distinct u32 {pct}% distinct"), |b| {
198+
b.iter(|| {
199+
let mut accumulator = prepare_accumulator(DataType::UInt32);
200+
accumulator
201+
.update_batch(std::slice::from_ref(&values))
202+
.unwrap()
203+
})
204+
});
205+
206+
// Int32
207+
let values = Arc::new(create_i32_array(n_distinct)) as ArrayRef;
208+
c.bench_function(&format!("count_distinct i32 {pct}% distinct"), |b| {
209+
b.iter(|| {
210+
let mut accumulator = prepare_accumulator(DataType::Int32);
211+
accumulator
212+
.update_batch(std::slice::from_ref(&values))
213+
.unwrap()
214+
})
215+
});
216+
}
175217
}
176218

177219
/// Create group indices with uniform distribution
@@ -204,10 +246,11 @@ fn count_distinct_groups_benchmark(c: &mut Criterion) {
204246
let cardinalities = [("low", 20), ("mid", 80), ("high", 99)];
205247
let distributions = ["uniform", "skewed"];
206248

249+
// i64 benchmarks
207250
for num_groups in group_counts {
208251
for (card_name, distinct_pct) in cardinalities {
209252
for dist in distributions {
210-
let name = format!("g{num_groups}_{card_name}_{dist}");
253+
let name = format!("i64_g{num_groups}_{card_name}_{dist}");
211254
let n_distinct = BATCH_SIZE * distinct_pct / 100;
212255
let values = Arc::new(create_i64_array(n_distinct)) as ArrayRef;
213256
let group_indices = if dist == "uniform" {
@@ -270,6 +313,142 @@ fn count_distinct_groups_benchmark(c: &mut Criterion) {
270313
}
271314
}
272315
}
316+
317+
// i32 benchmarks
318+
for num_groups in group_counts {
319+
for (card_name, distinct_pct) in cardinalities {
320+
for dist in distributions {
321+
let name = format!("i32_g{num_groups}_{card_name}_{dist}");
322+
let n_distinct = BATCH_SIZE * distinct_pct / 100;
323+
let values = Arc::new(create_i32_array(n_distinct)) as ArrayRef;
324+
let group_indices = if dist == "uniform" {
325+
create_uniform_groups(num_groups)
326+
} else {
327+
create_skewed_groups(num_groups)
328+
};
329+
330+
let (_schema, args) = prepare_args(DataType::Int32);
331+
332+
if count_fn.groups_accumulator_supported(args.clone()) {
333+
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
334+
b.iter(|| {
335+
let mut acc =
336+
count_fn.create_groups_accumulator(args.clone()).unwrap();
337+
acc.update_batch(
338+
std::slice::from_ref(&values),
339+
&group_indices,
340+
None,
341+
num_groups,
342+
)
343+
.unwrap();
344+
acc.evaluate(EmitTo::All).unwrap()
345+
})
346+
});
347+
} else {
348+
let arr = values.as_any().downcast_ref::<Int32Array>().unwrap();
349+
let mut group_rows: Vec<Vec<i32>> = vec![Vec::new(); num_groups];
350+
for (idx, &group_idx) in group_indices.iter().enumerate() {
351+
if arr.is_valid(idx) {
352+
group_rows[group_idx].push(arr.value(idx));
353+
}
354+
}
355+
let group_arrays: Vec<ArrayRef> = group_rows
356+
.iter()
357+
.map(|rows| Arc::new(Int32Array::from(rows.clone())) as ArrayRef)
358+
.collect();
359+
360+
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
361+
b.iter(|| {
362+
let mut accumulators: Vec<_> = (0..num_groups)
363+
.map(|_| prepare_accumulator(DataType::Int32))
364+
.collect();
365+
366+
for (group_idx, batch) in group_arrays.iter().enumerate() {
367+
if !batch.is_empty() {
368+
accumulators[group_idx]
369+
.update_batch(std::slice::from_ref(batch))
370+
.unwrap();
371+
}
372+
}
373+
374+
let _results: Vec<_> = accumulators
375+
.iter_mut()
376+
.map(|acc| acc.evaluate().unwrap())
377+
.collect();
378+
})
379+
});
380+
}
381+
}
382+
}
383+
}
384+
385+
// u32 benchmarks
386+
for num_groups in group_counts {
387+
for (card_name, distinct_pct) in cardinalities {
388+
for dist in distributions {
389+
let name = format!("u32_g{num_groups}_{card_name}_{dist}");
390+
let n_distinct = BATCH_SIZE * distinct_pct / 100;
391+
let values = Arc::new(create_u32_array(n_distinct)) as ArrayRef;
392+
let group_indices = if dist == "uniform" {
393+
create_uniform_groups(num_groups)
394+
} else {
395+
create_skewed_groups(num_groups)
396+
};
397+
398+
let (_schema, args) = prepare_args(DataType::UInt32);
399+
400+
if count_fn.groups_accumulator_supported(args.clone()) {
401+
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
402+
b.iter(|| {
403+
let mut acc =
404+
count_fn.create_groups_accumulator(args.clone()).unwrap();
405+
acc.update_batch(
406+
std::slice::from_ref(&values),
407+
&group_indices,
408+
None,
409+
num_groups,
410+
)
411+
.unwrap();
412+
acc.evaluate(EmitTo::All).unwrap()
413+
})
414+
});
415+
} else {
416+
let arr = values.as_any().downcast_ref::<UInt32Array>().unwrap();
417+
let mut group_rows: Vec<Vec<u32>> = vec![Vec::new(); num_groups];
418+
for (idx, &group_idx) in group_indices.iter().enumerate() {
419+
if arr.is_valid(idx) {
420+
group_rows[group_idx].push(arr.value(idx));
421+
}
422+
}
423+
let group_arrays: Vec<ArrayRef> = group_rows
424+
.iter()
425+
.map(|rows| Arc::new(UInt32Array::from(rows.clone())) as ArrayRef)
426+
.collect();
427+
428+
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
429+
b.iter(|| {
430+
let mut accumulators: Vec<_> = (0..num_groups)
431+
.map(|_| prepare_accumulator(DataType::UInt32))
432+
.collect();
433+
434+
for (group_idx, batch) in group_arrays.iter().enumerate() {
435+
if !batch.is_empty() {
436+
accumulators[group_idx]
437+
.update_batch(std::slice::from_ref(batch))
438+
.unwrap();
439+
}
440+
}
441+
442+
let _results: Vec<_> = accumulators
443+
.iter_mut()
444+
.map(|acc| acc.evaluate().unwrap())
445+
.collect();
446+
})
447+
});
448+
}
449+
}
450+
}
451+
}
273452
}
274453

275454
criterion_group!(

0 commit comments

Comments
 (0)