add group benches

Author: coderfender

datafusion/functions-aggregate/benches/count_distinct.rs

@@ -87,6 +87,13 @@ 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()));
@@ -174,101 +181,58 @@ fn count_distinct_benchmark(c: &mut Criterion) {
     });
 }
 
-/// Create group indices with uniform distribution
-fn create_uniform_groups(num_groups: usize) -> Vec<usize> {
-    let mut rng = StdRng::seed_from_u64(42);
-    (0..BATCH_SIZE)
-        .map(|_| rng.random_range(0..num_groups))
-        .collect()
-}
-
-/// Create group indices with skewed distribution (80% in 20% of groups)
-fn create_skewed_groups(num_groups: usize) -> Vec<usize> {
-    let mut rng = StdRng::seed_from_u64(42);
-    let hot_groups = (num_groups / 5).max(1);
-    (0..BATCH_SIZE)
-        .map(|_| {
-            if rng.random_range(0..100) < 80 {
-                rng.random_range(0..hot_groups)
-            } else {
-                rng.random_range(0..num_groups)
-            }
-        })
-        .collect()
-}
-
 fn count_distinct_groups_benchmark(c: &mut Criterion) {
     let count_fn = Count::new();
 
-    // bench different scenarios
-    let scenarios = [
-        // (name, num_groups, distinct_pct, group_fn)
-        ("sparse_uniform", 10, 80, "uniform"),
-        ("moderate_uniform", 100, 80, "uniform"),
-        ("dense_uniform", 1000, 80, "uniform"),
-        ("sparse_skewed", 10, 80, "skewed"),
-        ("dense_skewed", 1000, 80, "skewed"),
-        ("sparse_high_cardinality", 10, 99, "uniform"),
-        ("dense_low_cardinality", 1000, 20, "uniform"),
-    ];
-
-    for (name, num_groups, distinct_pct, group_type) in scenarios {
-        let n_distinct = BATCH_SIZE * distinct_pct / 100;
+    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 = if group_type == "uniform" {
-            create_uniform_groups(num_groups)
-        } else {
-            create_skewed_groups(num_groups)
-        };
+        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 {name}"), |b| {
-                b.iter(|| {
-                    let (_schema, args) = prepare_args(DataType::Int64);
-                    let mut acc = count_fn.create_groups_accumulator(args).unwrap();
-                    acc.update_batch(
-                        std::slice::from_ref(&values),
-                        &group_indices,
-                        None,
-                        num_groups,
-                    )
-                    .unwrap();
-                    acc.evaluate(EmitTo::All).unwrap()
-                })
-            });
+            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 {name}"), |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();
+            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();
-                })
-            });
+                        let _results: Vec<_> = accumulators
+                            .iter_mut()
+                            .map(|acc| acc.evaluate().unwrap())
+                            .collect();
+                    })
+                },
+            );
         }
     }
 }
 
-criterion_group!(
-    benches,
-    count_distinct_benchmark,
-    count_distinct_groups_benchmark
-);
+criterion_group!(benches, count_distinct_benchmark, count_distinct_groups_benchmark);
 criterion_main!(benches);