add group benches

Author: B Vadlamani

datafusion/functions-aggregate/benches/count_distinct.rs

@@ -23,7 +23,7 @@ use arrow::array::{
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{Criterion, criterion_group, criterion_main};
 use datafusion_expr::function::AccumulatorArgs;
-use datafusion_expr::{Accumulator, AggregateUDFImpl};
+use datafusion_expr::{Accumulator, AggregateUDFImpl, EmitTo};
 use datafusion_functions_aggregate::count::Count;
 use datafusion_physical_expr::expressions::col;
 use rand::rngs::StdRng;
@@ -87,6 +87,37 @@ fn create_i16_array(n_distinct: usize) -> Int16Array {
         .collect()
 }
 
+fn create_group_indices(num_groups: usize) -> Vec<usize> {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| rng.random_range(0..num_groups))
+        .collect()
+}
+
+fn prepare_args(data_type: DataType) -> (Arc<Schema>, AccumulatorArgs<'static>) {
+    let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)]));
+    let schema_leaked: &'static Schema = Box::leak(Box::new((*schema).clone()));
+    let expr = col("f", schema_leaked).unwrap();
+    let expr_leaked: &'static _ = Box::leak(Box::new(expr));
+    let return_field: Arc<Field> = Field::new("f", DataType::Int64, true).into();
+    let return_field_leaked: &'static _ = Box::leak(Box::new(return_field.clone()));
+    let expr_field = expr_leaked.return_field(schema_leaked).unwrap();
+    let expr_field_leaked: &'static _ = Box::leak(Box::new(expr_field));
+
+    let accumulator_args = AccumulatorArgs {
+        return_field: return_field_leaked.clone(),
+        schema: schema_leaked,
+        expr_fields: std::slice::from_ref(expr_field_leaked),
+        ignore_nulls: false,
+        order_bys: &[],
+        is_reversed: false,
+        name: "count(distinct f)",
+        is_distinct: true,
+        exprs: std::slice::from_ref(expr_leaked),
+    };
+    (schema, accumulator_args)
+}
+
 fn count_distinct_benchmark(c: &mut Criterion) {
     for pct in [80, 99] {
         let n_distinct = BATCH_SIZE * pct / 100;
@@ -150,5 +181,58 @@ fn count_distinct_benchmark(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, count_distinct_benchmark);
+fn count_distinct_groups_benchmark(c: &mut Criterion) {
+    let count_fn = Count::new();
+
+    for num_groups in [10, 100, 1000] {
+        let n_distinct = BATCH_SIZE * 80 / 100;
+        let values = Arc::new(create_i64_array(n_distinct)) as ArrayRef;
+        let group_indices = create_group_indices(num_groups);
+
+        let (_schema, args) = prepare_args(DataType::Int64);
+
+        if count_fn.groups_accumulator_supported(args.clone()) {
+            c.bench_function(
+                &format!("count_distinct_groups i64 {num_groups} groups (GroupsAccumulator)"),
+                |b| {
+                    b.iter(|| {
+                        let (_schema, args) = prepare_args(DataType::Int64);
+                        let mut acc = count_fn.create_groups_accumulator(args).unwrap();
+                        acc.update_batch(&[values.clone()], &group_indices, None, num_groups)
+                            .unwrap();
+                        acc.evaluate(EmitTo::All).unwrap()
+                    })
+                },
+            );
+        } else {
+            c.bench_function(
+                &format!("count_distinct_groups i64 {num_groups} groups (N Accumulators)"),
+                |b| {
+                    b.iter(|| {
+                        let mut accumulators: Vec<_> = (0..num_groups)
+                            .map(|_| prepare_accumulator(DataType::Int64))
+                            .collect();
+
+                        let arr = values.as_any().downcast_ref::<Int64Array>().unwrap();
+                        for (idx, group_idx) in group_indices.iter().enumerate() {
+                            if let Some(val) = arr.value(idx).into() {
+                                let single_val = Arc::new(Int64Array::from(vec![Some(val)])) as ArrayRef;
+                                accumulators[*group_idx]
+                                    .update_batch(std::slice::from_ref(&single_val))
+                                    .unwrap();
+                            }
+                        }
+
+                        let _results: Vec<_> = accumulators
+                            .iter_mut()
+                            .map(|acc| acc.evaluate().unwrap())
+                            .collect();
+                    })
+                },
+            );
+        }
+    }
+}
+
+criterion_group!(benches, count_distinct_benchmark, count_distinct_groups_benchmark);
 criterion_main!(benches);