add_32_bit_ints

Author: B Vadlamani

datafusion/functions-aggregate/benches/count_distinct.rs

@@ -18,7 +18,8 @@
 use std::sync::Arc;
 
 use arrow::array::{
-    Array, ArrayRef, Int8Array, Int16Array, Int64Array, UInt8Array, UInt16Array,
+    ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array, UInt8Array, UInt16Array,
+    UInt32Array,
 };
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{Criterion, criterion_group, criterion_main};
@@ -87,6 +88,20 @@ fn create_i16_array(n_distinct: usize) -> Int16Array {
         .collect()
 }
 
+fn create_u32_array(n_distinct: usize) -> UInt32Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..n_distinct as u32)))
+        .collect()
+}
+
+fn create_i32_array(n_distinct: usize) -> Int32Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..n_distinct as i32)))
+        .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()));
@@ -172,6 +187,36 @@ fn count_distinct_benchmark(c: &mut Criterion) {
                 .unwrap()
         })
     });
+
+    // 32-bit integer types
+
+    // UInt32
+    for pct in [80, 99] {
+        let n_distinct = BATCH_SIZE * pct / 100;
+        let values = Arc::new(create_u32_array(n_distinct)) as ArrayRef;
+        c.bench_function(&format!("count_distinct u32 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::UInt32);
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+    }
+
+    // Int32
+    for pct in [80, 99] {
+        let n_distinct = BATCH_SIZE * pct / 100;
+        let values = Arc::new(create_i32_array(n_distinct)) as ArrayRef;
+        c.bench_function(&format!("count_distinct i32 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Int32);
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+    }
 }
 
 /// Create group indices with uniform distribution
@@ -197,75 +242,66 @@ fn create_skewed_groups(num_groups: usize) -> Vec<usize> {
         .collect()
 }
 
+fn create_array_for_type(data_type: &DataType, n_distinct: usize) -> ArrayRef {
+    match data_type {
+        DataType::Int64 => Arc::new(create_i64_array(n_distinct)) as ArrayRef,
+        DataType::Int32 => Arc::new(create_i32_array(n_distinct)) as ArrayRef,
+        DataType::UInt32 => Arc::new(create_u32_array(n_distinct)) as ArrayRef,
+        DataType::Int16 => Arc::new(create_i16_array(n_distinct)) as ArrayRef,
+        DataType::UInt16 => Arc::new(create_u16_array(n_distinct)) as ArrayRef,
+        DataType::Int8 => Arc::new(create_i8_array(n_distinct)) as ArrayRef,
+        DataType::UInt8 => Arc::new(create_u8_array(n_distinct)) as ArrayRef,
+        _ => panic!("Unsupported type for benchmark"),
+    }
+}
+
 fn count_distinct_groups_benchmark(c: &mut Criterion) {
     let count_fn = Count::new();
 
     let group_counts = [100, 1000, 10000];
     let cardinalities = [("low", 20), ("mid", 80), ("high", 99)];
     let distributions = ["uniform", "skewed"];
+    let data_types = [
+        ("i64", DataType::Int64),
+        ("u32", DataType::UInt32),
+        ("i32", DataType::Int32),
+        ("u16", DataType::UInt16),
+        ("i16", DataType::Int16),
+        ("u8", DataType::UInt8),
+        ("i8", DataType::Int8),
+    ];
+
+    for (type_name, data_type) in data_types {
+        for num_groups in group_counts {
+            for (card_name, distinct_pct) in cardinalities {
+                for dist in distributions {
+                    let name = format!("{type_name}_g{num_groups}_{card_name}_{dist}");
+                    let n_distinct = BATCH_SIZE * distinct_pct / 100;
+                    let values = create_array_for_type(&data_type, n_distinct);
+                    let group_indices = if dist == "uniform" {
+                        create_uniform_groups(num_groups)
+                    } else {
+                        create_skewed_groups(num_groups)
+                    };
+
+                    let (_schema, args) = prepare_args(data_type.clone());
 
-    for num_groups in group_counts {
-        for (card_name, distinct_pct) in cardinalities {
-            for dist in distributions {
-                let name = format!("g{num_groups}_{card_name}_{dist}");
-                let n_distinct = BATCH_SIZE * distinct_pct / 100;
-                let values = Arc::new(create_i64_array(n_distinct)) as ArrayRef;
-                let group_indices = if dist == "uniform" {
-                    create_uniform_groups(num_groups)
-                } else {
-                    create_skewed_groups(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 mut acc =
-                                count_fn.create_groups_accumulator(args.clone()).unwrap();
-                            acc.update_batch(
-                                std::slice::from_ref(&values),
-                                &group_indices,
-                                None,
-                                num_groups,
-                            )
-                            .unwrap();
-                            acc.evaluate(EmitTo::All).unwrap()
-                        })
-                    });
-                } else {
-                    let arr = values.as_any().downcast_ref::<Int64Array>().unwrap();
-                    let mut group_rows: Vec<Vec<i64>> = vec![Vec::new(); num_groups];
-                    for (idx, &group_idx) in group_indices.iter().enumerate() {
-                        if arr.is_valid(idx) {
-                            group_rows[group_idx].push(arr.value(idx));
-                        }
+                    if count_fn.groups_accumulator_supported(args.clone()) {
+                        c.bench_function(&format!("count_distinct_groups {name}"), |b| {
+                            b.iter(|| {
+                                let mut acc =
+                                    count_fn.create_groups_accumulator(args.clone()).unwrap();
+                                acc.update_batch(
+                                    std::slice::from_ref(&values),
+                                    &group_indices,
+                                    None,
+                                    num_groups,
+                                )
+                                .unwrap();
+                                acc.evaluate(EmitTo::All).unwrap()
+                            })
+                        });
                     }
-                    let group_arrays: Vec<ArrayRef> = group_rows
-                        .iter()
-                        .map(|rows| Arc::new(Int64Array::from(rows.clone())) as ArrayRef)
-                        .collect();
-
-                    c.bench_function(&format!("count_distinct_groups {name}"), |b| {
-                        b.iter(|| {
-                            let mut accumulators: Vec<_> = (0..num_groups)
-                                .map(|_| prepare_accumulator(DataType::Int64))
-                                .collect();
-
-                            for (group_idx, batch) in group_arrays.iter().enumerate() {
-                                if !batch.is_empty() {
-                                    accumulators[group_idx]
-                                        .update_batch(std::slice::from_ref(batch))
-                                        .unwrap();
-                                }
-                            }
-
-                            let _results: Vec<_> = accumulators
-                                .iter_mut()
-                                .map(|acc| acc.evaluate().unwrap())
-                                .collect();
-                        })
-                    });
                 }
             }
         }