add more bench cases

arrow/benches/partition_kernels.rs

@@ -23,71 +23,119 @@ extern crate arrow;
 use arrow::compute::kernels::partition::lexicographical_partition_ranges;
 use arrow::compute::kernels::sort::{lexsort, SortColumn};
 use arrow::util::bench_util::*;
-use arrow::{array::*, datatypes::UInt8Type};
+use arrow::{
+    array::*,
+    datatypes::{ArrowPrimitiveType, Float64Type, UInt8Type},
+};
+use rand::distributions::{Distribution, Standard};
+use std::iter;
 
-fn create_array(size: usize, with_nulls: bool) -> ArrayRef {
+fn create_array<T: ArrowPrimitiveType>(size: usize, with_nulls: bool) -> ArrayRef
+where
+    Standard: Distribution<T::Native>,
+{
     let null_density = if with_nulls { 0.5 } else { 0.0 };
-    let array = create_primitive_array::<UInt8Type>(size, null_density);
+    let array = create_primitive_array::<T>(size, null_density);
     Arc::new(array)
 }
 
-fn bench_partition(array_a: &ArrayRef, array_b: &ArrayRef) {
-    let columns = vec![
-        SortColumn {
-            values: array_a.clone(),
+fn bench_partition(sorted_columns: &[ArrayRef]) {
+    let columns = sorted_columns
+        .iter()
+        .map(|arr| SortColumn {
+            values: arr.clone(),
             options: None,
-        },
-        SortColumn {
-            values: array_b.clone(),
-            options: None,
-        },
-    ];
+        })
+        .collect::<Vec<_>>();
 
     criterion::black_box(lexicographical_partition_ranges(&columns).unwrap());
 }
 
-fn create_sorted_data(pow: u32, with_nulls: bool) -> (ArrayRef, ArrayRef) {
-    let sorted = lexsort(
+fn create_sorted_low_cardinality_data(length: usize) -> Vec<ArrayRef> {
+    let arr = Int64Array::from_iter_values(
+        iter::repeat(1)
+            .take(length / 4)
+            .chain(iter::repeat(2).take(length / 4))
+            .chain(iter::repeat(3).take(length / 4))
+            .chain(iter::repeat(4).take(length / 4)),
+    );
+    lexsort(
+        &[SortColumn {
+            values: Arc::new(arr),
+            options: None,
+        }],
+        None,
+    )
+    .unwrap()
+}
+
+fn create_sorted_float_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> {
+    lexsort(
         &[
             SortColumn {
-                values: create_array(2u64.pow(pow) as usize, with_nulls),
+                values: create_array::<Float64Type>(2u64.pow(pow) as usize, with_nulls),
                 options: None,
             },
             SortColumn {
-                values: create_array(2u64.pow(pow) as usize, with_nulls),
+                values: create_array::<Float64Type>(2u64.pow(pow) as usize, with_nulls),
                 options: None,
             },
         ],
         None,
     )
-    .unwrap();
+    .unwrap()
+}
 
-    // must pre-sort the result in order to have the expected behavior
-    let arr_a = sorted[0].clone();
-    let arr_b = sorted[1].clone();
-    (arr_a, arr_b)
+fn create_sorted_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> {
+    lexsort(
+        &[
+            SortColumn {
+                values: create_array::<UInt8Type>(2u64.pow(pow) as usize, with_nulls),
+                options: None,
+            },
+            SortColumn {
+                values: create_array::<UInt8Type>(2u64.pow(pow) as usize, with_nulls),
+                options: None,
+            },
+        ],
+        None,
+    )
+    .unwrap()
 }
 
 fn add_benchmark(c: &mut Criterion) {
-    let (arr_a, arr_b) = create_sorted_data(10, false);
-    c.bench_function("partition 2^10", |b| {
-        b.iter(|| bench_partition(&arr_a, &arr_b))
+    let sorted_columns = create_sorted_data(10, false);
+    c.bench_function("lexicographical_partition_ranges(u8) 2^10", |b| {
+        b.iter(|| bench_partition(&sorted_columns))
     });
 
-    let (arr_a, arr_b) = create_sorted_data(12, false);
-    c.bench_function("partition 2^12", |b| {
-        b.iter(|| bench_partition(&arr_a, &arr_b))
+    let sorted_columns = create_sorted_data(12, false);
+    c.bench_function("lexicographical_partition_ranges(u8) 2^12", |b| {
+        b.iter(|| bench_partition(&sorted_columns))
     });
 
-    let (arr_a, arr_b) = create_sorted_data(10, true);
-    c.bench_function("partition 2^10 with nulls", |b| {
-        b.iter(|| bench_partition(&arr_a, &arr_b))
-    });
+    let sorted_columns = create_sorted_data(10, true);
+    c.bench_function(
+        "lexicographical_partition_ranges(u8) 2^10 with nulls",
+        |b| b.iter(|| bench_partition(&sorted_columns)),
+    );
 
-    let (arr_a, arr_b) = create_sorted_data(12, true);
-    c.bench_function("partition 2^12 with nulls", |b| {
-        b.iter(|| bench_partition(&arr_a, &arr_b))
+    let sorted_columns = create_sorted_data(12, true);
+    c.bench_function(
+        "lexicographical_partition_ranges(u8) 2^12 with nulls",
+        |b| b.iter(|| bench_partition(&sorted_columns)),
+    );
+
+    let sorted_columns = create_sorted_float_data(10, false);
+    c.bench_function("lexicographical_partition_ranges(f64) 2^10", |b| {
+        b.iter(|| bench_partition(&sorted_columns))
     });
+
+    let sorted_columns = create_sorted_low_cardinality_data(1024);
+    c.bench_function(
+        "lexicographical_partition_ranges(low cardinality) 1024",
+        |b| b.iter(|| bench_partition(&sorted_columns)),
+    );
 }
 
 criterion_group!(benches, add_benchmark);