create_benches_small_int

Author: coderfender

datafusion/functions-aggregate/benches/approx_distinct.rs

@@ -17,7 +17,10 @@
 
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, Int64Array, StringArray, StringViewArray};
+use arrow::array::{
+    ArrayRef, Int8Array, Int16Array, Int64Array, StringArray, StringViewArray,
+    UInt8Array, UInt16Array,
+};
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{Criterion, criterion_group, criterion_main};
 use datafusion_expr::function::AccumulatorArgs;
@@ -56,6 +59,38 @@ fn create_i64_array(n_distinct: usize) -> Int64Array {
         .collect()
 }
 
+fn create_u8_array(n_distinct: usize) -> UInt8Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    let max_val = n_distinct.min(256) as u8;
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..max_val)))
+        .collect()
+}
+
+fn create_i8_array(n_distinct: usize) -> Int8Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    let max_val = (n_distinct.min(256) / 2) as i8;
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(-max_val..max_val)))
+        .collect()
+}
+
+fn create_u16_array(n_distinct: usize) -> UInt16Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    let max_val = n_distinct.min(65536) as u16;
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..max_val)))
+        .collect()
+}
+
+fn create_i16_array(n_distinct: usize) -> Int16Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    let max_val = (n_distinct.min(65536) / 2) as i16;
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(-max_val..max_val)))
+        .collect()
+}
+
 /// Creates a pool of `n_distinct` random strings of the given length.
 fn create_string_pool(n_distinct: usize, string_length: usize) -> Vec<String> {
     let mut rng = StdRng::seed_from_u64(42);
@@ -133,6 +168,52 @@ fn approx_distinct_benchmark(c: &mut Criterion) {
             );
         }
     }
+
+    // Small integer types
+
+    // UInt8
+    let values = Arc::new(create_u8_array(200)) as ArrayRef;
+    c.bench_function("approx_distinct u8 bitmap", |b| {
+        b.iter(|| {
+            let mut accumulator = prepare_accumulator(DataType::UInt8);
+            accumulator
+                .update_batch(std::slice::from_ref(&values))
+                .unwrap()
+        })
+    });
+
+    // Int8
+    let values = Arc::new(create_i8_array(200)) as ArrayRef;
+    c.bench_function("approx_distinct i8 bitmap", |b| {
+        b.iter(|| {
+            let mut accumulator = prepare_accumulator(DataType::Int8);
+            accumulator
+                .update_batch(std::slice::from_ref(&values))
+                .unwrap()
+        })
+    });
+
+    // UInt16
+    let values = Arc::new(create_u16_array(50000)) as ArrayRef;
+    c.bench_function("approx_distinct u16 bitmap", |b| {
+        b.iter(|| {
+            let mut accumulator = prepare_accumulator(DataType::UInt16);
+            accumulator
+                .update_batch(std::slice::from_ref(&values))
+                .unwrap()
+        })
+    });
+
+    // Int16
+    let values = Arc::new(create_i16_array(50000)) as ArrayRef;
+    c.bench_function("approx_distinct i16 bitmap", |b| {
+        b.iter(|| {
+            let mut accumulator = prepare_accumulator(DataType::Int16);
+            accumulator
+                .update_batch(std::slice::from_ref(&values))
+                .unwrap()
+        })
+    });
 }
 
 criterion_group!(benches, approx_distinct_benchmark);