perf: add to aggregate_vectorized PrimitiveGroupValueBuilder benchmark as well (#17930)

I have pretty major perf optimization to primitive impl so adding this before I open my PR

Author: rluvaton

datafusion/physical-plan/benches/aggregate_vectorized.rs

@@ -16,19 +16,30 @@
 // under the License.
 
 use arrow::array::ArrayRef;
-use arrow::datatypes::StringViewType;
+use arrow::datatypes::{Int32Type, StringViewType};
 use arrow::util::bench_util::{
-    create_string_view_array_with_len, create_string_view_array_with_max_len,
+    create_primitive_array, create_string_view_array_with_len,
+    create_string_view_array_with_max_len,
+};
+use arrow_schema::DataType;
+use criterion::measurement::WallTime;
+use criterion::{
+    criterion_group, criterion_main, BenchmarkGroup, BenchmarkId, Criterion,
 };
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use datafusion_physical_plan::aggregates::group_values::multi_group_by::bytes_view::ByteViewGroupValueBuilder;
+use datafusion_physical_plan::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder;
 use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupColumn;
 use std::sync::Arc;
 
 const SIZES: [usize; 3] = [1_000, 10_000, 100_000];
 const NULL_DENSITIES: [f32; 3] = [0.0, 0.1, 0.5];
 
 fn bench_vectorized_append(c: &mut Criterion) {
+    byte_view_vectorized_append(c);
+    primitive_vectorized_append(c);
+}
+
+fn byte_view_vectorized_append(c: &mut Criterion) {
     let mut group = c.benchmark_group("ByteViewGroupValueBuilder_vectorized_append");
 
     for &size in &SIZES {
@@ -38,150 +49,153 @@ fn bench_vectorized_append(c: &mut Criterion) {
             let input = create_string_view_array_with_len(size, null_density, 8, false);
             let input: ArrayRef = Arc::new(input);
 
-            // vectorized_append
-            let id = BenchmarkId::new(
-                format!("inlined_null_{null_density:.1}_size_{size}"),
-                "vectorized_append",
-            );
-            group.bench_function(id, |b| {
-                b.iter(|| {
-                    let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                    builder.vectorized_append(&input, &rows).unwrap();
-                });
-            });
-
-            // append_val
-            let id = BenchmarkId::new(
-                format!("inlined_null_{null_density:.1}_size_{size}"),
-                "append_val",
-            );
-            group.bench_function(id, |b| {
-                b.iter(|| {
-                    let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                    for &i in &rows {
-                        builder.append_val(&input, i).unwrap();
-                    }
-                });
-            });
-
-            // vectorized_equal_to
-            let id = BenchmarkId::new(
-                format!("inlined_null_{null_density:.1}_size_{size}"),
-                "vectorized_equal_to",
-            );
-            group.bench_function(id, |b| {
-                let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                builder.vectorized_append(&input, &rows).unwrap();
-                let mut results = vec![true; size];
-                b.iter(|| {
-                    builder.vectorized_equal_to(&rows, &input, &rows, &mut results);
-                });
-            });
+            bytes_bench(&mut group, "inline", size, &rows, null_density, &input);
         }
     }
 
     for &size in &SIZES {
         let rows: Vec<usize> = (0..size).collect();
 
         for &null_density in &NULL_DENSITIES {
-            let scenario = "mixed";
             let input = create_string_view_array_with_len(size, null_density, 64, true);
             let input: ArrayRef = Arc::new(input);
 
-            // vectorized_append
-            let id = BenchmarkId::new(
-                format!("{scenario}_null_{null_density:.1}_size_{size}"),
-                "vectorized_append",
-            );
-            group.bench_function(id, |b| {
-                b.iter(|| {
-                    let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                    builder.vectorized_append(&input, &rows).unwrap();
-                });
-            });
-
-            // append_val
-            let id = BenchmarkId::new(
-                format!("{scenario}_null_{null_density:.1}_size_{size}"),
-                "append_val",
-            );
-            group.bench_function(id, |b| {
-                b.iter(|| {
-                    let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                    for &i in &rows {
-                        builder.append_val(&input, i).unwrap();
-                    }
-                });
-            });
-
-            // vectorized_equal_to
-            let id = BenchmarkId::new(
-                format!("{scenario}_null_{null_density:.1}_size_{size}"),
-                "vectorized_equal_to",
-            );
-            group.bench_function(id, |b| {
-                let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                builder.vectorized_append(&input, &rows).unwrap();
-                let mut results = vec![true; size];
-                b.iter(|| {
-                    builder.vectorized_equal_to(&rows, &input, &rows, &mut results);
-                });
-            });
+            bytes_bench(&mut group, "scenario", size, &rows, null_density, &input);
         }
     }
 
     for &size in &SIZES {
         let rows: Vec<usize> = (0..size).collect();
 
         for &null_density in &NULL_DENSITIES {
-            let scenario = "random";
             let input = create_string_view_array_with_max_len(size, null_density, 400);
             let input: ArrayRef = Arc::new(input);
 
-            // vectorized_append
-            let id = BenchmarkId::new(
-                format!("{scenario}_null_{null_density:.1}_size_{size}"),
-                "vectorized_append",
-            );
-            group.bench_function(id, |b| {
-                b.iter(|| {
-                    let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                    builder.vectorized_append(&input, &rows).unwrap();
-                });
-            });
-
-            // append_val
-            let id = BenchmarkId::new(
-                format!("{scenario}_null_{null_density:.1}_size_{size}"),
-                "append_val",
-            );
-            group.bench_function(id, |b| {
-                b.iter(|| {
-                    let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                    for &i in &rows {
-                        builder.append_val(&input, i).unwrap();
-                    }
-                });
-            });
-
-            // vectorized_equal_to
-            let id = BenchmarkId::new(
-                format!("{scenario}_null_{null_density:.1}_size_{size}"),
-                "vectorized_equal_to",
-            );
-            group.bench_function(id, |b| {
-                let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
-                builder.vectorized_append(&input, &rows).unwrap();
-                let mut results = vec![true; size];
-                b.iter(|| {
-                    builder.vectorized_equal_to(&rows, &input, &rows, &mut results);
-                });
-            });
+            bytes_bench(&mut group, "random", size, &rows, null_density, &input);
+        }
+    }
+
+    group.finish();
+}
+
+fn bytes_bench(
+    group: &mut BenchmarkGroup<WallTime>,
+    bench_prefix: &str,
+    size: usize,
+    rows: &Vec<usize>,
+    null_density: f32,
+    input: &ArrayRef,
+) {
+    // vectorized_append
+    let id = BenchmarkId::new(
+        format!("{bench_prefix}_null_{null_density:.1}_size_{size}"),
+        "vectorized_append",
+    );
+    group.bench_function(id, |b| {
+        b.iter(|| {
+            let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
+            builder.vectorized_append(input, rows).unwrap();
+        });
+    });
+
+    // append_val
+    let id = BenchmarkId::new(
+        format!("{bench_prefix}_null_{null_density:.1}_size_{size}"),
+        "append_val",
+    );
+    group.bench_function(id, |b| {
+        b.iter(|| {
+            let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
+            for &i in rows {
+                builder.append_val(input, i).unwrap();
+            }
+        });
+    });
+
+    // vectorized_equal_to
+    let id = BenchmarkId::new(
+        format!("{bench_prefix}_null_{null_density:.1}_size_{size}"),
+        "vectorized_equal_to",
+    );
+    group.bench_function(id, |b| {
+        let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new();
+        builder.vectorized_append(input, rows).unwrap();
+        let mut results = vec![true; size];
+        b.iter(|| {
+            builder.vectorized_equal_to(rows, input, rows, &mut results);
+        });
+    });
+}
+
+fn primitive_vectorized_append(c: &mut Criterion) {
+    let mut group = c.benchmark_group("PrimitiveGroupValueBuilder_vectorized_append");
+
+    for &size in &SIZES {
+        let rows: Vec<usize> = (0..size).collect();
+
+        for &null_density in &NULL_DENSITIES {
+            if null_density == 0.0 {
+                bench_single_primitive::<false>(&mut group, size, &rows, null_density)
+            }
+            bench_single_primitive::<true>(&mut group, size, &rows, null_density);
         }
     }
 
     group.finish();
 }
 
+fn bench_single_primitive<const NULLABLE: bool>(
+    group: &mut BenchmarkGroup<WallTime>,
+    size: usize,
+    rows: &Vec<usize>,
+    null_density: f32,
+) {
+    if !NULLABLE {
+        assert_eq!(
+            null_density, 0.0,
+            "non-nullable case must have null_density 0"
+        );
+    }
+
+    let input = create_primitive_array::<Int32Type>(size, null_density);
+    let input: ArrayRef = Arc::new(input);
+    let function_name = format!("null_{null_density:.1}_nullable_{NULLABLE}_size_{size}");
+
+    // vectorized_append
+    let id = BenchmarkId::new(&function_name, "vectorized_append");
+    group.bench_function(id, |b| {
+        b.iter(|| {
+            let mut builder =
+                PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32);
+            builder.vectorized_append(&input, rows).unwrap();
+        });
+    });
+
+    // append_val
+    let id = BenchmarkId::new(&function_name, "append_val");
+    group.bench_function(id, |b| {
+        b.iter(|| {
+            let mut builder =
+                PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32);
+            for &i in rows {
+                builder.append_val(&input, i).unwrap();
+            }
+        });
+    });
+
+    // vectorized_equal_to
+    let id = BenchmarkId::new(&function_name, "vectorized_equal_to");
+    group.bench_function(id, |b| {
+        let mut builder =
+            PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32);
+        builder.vectorized_append(&input, rows).unwrap();
+        let mut results = vec![true; size];
+
+        b.iter(|| {
+            builder.vectorized_equal_to(rows, &input, rows, &mut results);
+        });
+    });
+}
+
 criterion_group!(benches, bench_vectorized_append);
 criterion_main!(benches);

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs

@@ -20,7 +20,7 @@
 mod boolean;
 mod bytes;
 pub mod bytes_view;
-mod primitive;
+pub mod primitive;
 
 use std::mem::{self, size_of};