Skip to content

Commit 1ff3d88

Browse files
author
B Vadlamani
committed
add int 32 bit data types benches
1 parent 5ccb2ac commit 1ff3d88

1 file changed

Lines changed: 204 additions & 61 deletions

File tree

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

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

2020
use arrow::array::{
21-
ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array, UInt8Array, UInt16Array,
22-
UInt32Array,
21+
Array, ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array, UInt8Array,
22+
UInt16Array, UInt32Array,
2323
};
2424
use arrow::datatypes::{DataType, Field, Schema};
2525
use criterion::{Criterion, criterion_group, criterion_main};
@@ -189,10 +189,10 @@ fn count_distinct_benchmark(c: &mut Criterion) {
189189
});
190190

191191
// 32-bit integer types
192-
193-
// UInt32
194192
for pct in [80, 99] {
195193
let n_distinct = BATCH_SIZE * pct / 100;
194+
195+
// UInt32
196196
let values = Arc::new(create_u32_array(n_distinct)) as ArrayRef;
197197
c.bench_function(&format!("count_distinct u32 {pct}% distinct"), |b| {
198198
b.iter(|| {
@@ -202,11 +202,8 @@ fn count_distinct_benchmark(c: &mut Criterion) {
202202
.unwrap()
203203
})
204204
});
205-
}
206205

207-
// Int32
208-
for pct in [80, 99] {
209-
let n_distinct = BATCH_SIZE * pct / 100;
206+
// Int32
210207
let values = Arc::new(create_i32_array(n_distinct)) as ArrayRef;
211208
c.bench_function(&format!("count_distinct i32 {pct}% distinct"), |b| {
212209
b.iter(|| {
@@ -242,66 +239,212 @@ fn create_skewed_groups(num_groups: usize) -> Vec<usize> {
242239
.collect()
243240
}
244241

245-
fn create_array_for_type(data_type: &DataType, n_distinct: usize) -> ArrayRef {
246-
match data_type {
247-
DataType::Int64 => Arc::new(create_i64_array(n_distinct)) as ArrayRef,
248-
DataType::Int32 => Arc::new(create_i32_array(n_distinct)) as ArrayRef,
249-
DataType::UInt32 => Arc::new(create_u32_array(n_distinct)) as ArrayRef,
250-
DataType::Int16 => Arc::new(create_i16_array(n_distinct)) as ArrayRef,
251-
DataType::UInt16 => Arc::new(create_u16_array(n_distinct)) as ArrayRef,
252-
DataType::Int8 => Arc::new(create_i8_array(n_distinct)) as ArrayRef,
253-
DataType::UInt8 => Arc::new(create_u8_array(n_distinct)) as ArrayRef,
254-
_ => panic!("Unsupported type for benchmark"),
255-
}
256-
}
257-
258242
fn count_distinct_groups_benchmark(c: &mut Criterion) {
259243
let count_fn = Count::new();
260244

261245
let group_counts = [100, 1000, 10000];
262246
let cardinalities = [("low", 20), ("mid", 80), ("high", 99)];
263247
let distributions = ["uniform", "skewed"];
264-
let data_types = [
265-
("i64", DataType::Int64),
266-
("u32", DataType::UInt32),
267-
("i32", DataType::Int32),
268-
("u16", DataType::UInt16),
269-
("i16", DataType::Int16),
270-
("u8", DataType::UInt8),
271-
("i8", DataType::Int8),
272-
];
273-
274-
for (type_name, data_type) in data_types {
275-
for num_groups in group_counts {
276-
for (card_name, distinct_pct) in cardinalities {
277-
for dist in distributions {
278-
let name = format!("{type_name}_g{num_groups}_{card_name}_{dist}");
279-
let n_distinct = BATCH_SIZE * distinct_pct / 100;
280-
let values = create_array_for_type(&data_type, n_distinct);
281-
let group_indices = if dist == "uniform" {
282-
create_uniform_groups(num_groups)
283-
} else {
284-
create_skewed_groups(num_groups)
285-
};
286-
287-
let (_schema, args) = prepare_args(data_type.clone());
288-
289-
if count_fn.groups_accumulator_supported(args.clone()) {
290-
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
291-
b.iter(|| {
292-
let mut acc =
293-
count_fn.create_groups_accumulator(args.clone()).unwrap();
294-
acc.update_batch(
295-
std::slice::from_ref(&values),
296-
&group_indices,
297-
None,
298-
num_groups,
299-
)
300-
.unwrap();
301-
acc.evaluate(EmitTo::All).unwrap()
302-
})
303-
});
248+
249+
// i64 benchmarks
250+
for num_groups in group_counts {
251+
for (card_name, distinct_pct) in cardinalities {
252+
for dist in distributions {
253+
let name = format!("i64_g{num_groups}_{card_name}_{dist}");
254+
let n_distinct = BATCH_SIZE * distinct_pct / 100;
255+
let values = Arc::new(create_i64_array(n_distinct)) as ArrayRef;
256+
let group_indices = if dist == "uniform" {
257+
create_uniform_groups(num_groups)
258+
} else {
259+
create_skewed_groups(num_groups)
260+
};
261+
262+
let (_schema, args) = prepare_args(DataType::Int64);
263+
264+
if count_fn.groups_accumulator_supported(args.clone()) {
265+
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
266+
b.iter(|| {
267+
let mut acc =
268+
count_fn.create_groups_accumulator(args.clone()).unwrap();
269+
acc.update_batch(
270+
std::slice::from_ref(&values),
271+
&group_indices,
272+
None,
273+
num_groups,
274+
)
275+
.unwrap();
276+
acc.evaluate(EmitTo::All).unwrap()
277+
})
278+
});
279+
} else {
280+
let arr = values.as_any().downcast_ref::<Int64Array>().unwrap();
281+
let mut group_rows: Vec<Vec<i64>> = vec![Vec::new(); num_groups];
282+
for (idx, &group_idx) in group_indices.iter().enumerate() {
283+
if arr.is_valid(idx) {
284+
group_rows[group_idx].push(arr.value(idx));
285+
}
286+
}
287+
let group_arrays: Vec<ArrayRef> = group_rows
288+
.iter()
289+
.map(|rows| Arc::new(Int64Array::from(rows.clone())) as ArrayRef)
290+
.collect();
291+
292+
c.bench_function(&format!("count_distinct_groups {name}"), |b| {
293+
b.iter(|| {
294+
let mut accumulators: Vec<_> = (0..num_groups)
295+
.map(|_| prepare_accumulator(DataType::Int64))
296+
.collect();
297+
298+
for (group_idx, batch) in group_arrays.iter().enumerate() {
299+
if !batch.is_empty() {
300+
accumulators[group_idx]
301+
.update_batch(std::slice::from_ref(batch))
302+
.unwrap();
303+
}
304+
}
305+
306+
let _results: Vec<_> = accumulators
307+
.iter_mut()
308+
.map(|acc| acc.evaluate().unwrap())
309+
.collect();
310+
})
311+
});
312+
}
313+
}
314+
}
315+
}
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+
}
304422
}
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+
});
305448
}
306449
}
307450
}

0 commit comments

Comments
 (0)