Add List and ListView take benchmarks

arrow/benches/take_kernels.rs

@@ -186,6 +186,66 @@ fn add_benchmark(c: &mut Criterion) {
         b.iter(|| bench_take(&values, &indices))
     });
 
+    let values = create_primitive_list_array::<i32, Int32Type>(512, 0.0, 0.0, 20);
+    let indices = create_random_index(512, 0.0);
+    c.bench_function("take list i32 512", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.0, 0.0, 20);
+    let indices = create_random_index(1024, 0.0);
+    c.bench_function("take list i32 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.5, 0.0, 20);
+    let indices = create_random_index(1024, 0.0);
+    c.bench_function("take list i32 null values 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.0, 0.0, 202);
+    let indices = create_random_index(1024, 0.5);
+    c.bench_function("take list i32 null indices 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.5, 0.5, 20);
+    let indices = create_random_index(1024, 0.5);
+    c.bench_function("take list i32 null values null indices 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_view_array::<i32, Int32Type>(512, 0.0, 0.0, 20);
+    let indices = create_random_index(512, 0.0);
+    c.bench_function("take listview i32 512", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.0, 0.0, 20);
+    let indices = create_random_index(1024, 0.0);
+    c.bench_function("take listview i32 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.5, 0.0, 20);
+    let indices = create_random_index(1024, 0.0);
+    c.bench_function("take listview i32 null values 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.0, 0.0, 20);
+    let indices = create_random_index(1024, 0.5);
+    c.bench_function("take listview i32 null indices 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
+    let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.5, 0.5, 20);
+    let indices = create_random_index(1024, 0.5);
+    c.bench_function("take listview i32 null values null indices 1024", |b| {
+        b.iter(|| bench_take(&values, &indices))
+    });
+
     let values = create_primitive_run_array::<Int32Type, Int32Type>(1024, 512);
     let indices = create_random_index(1024, 0.0);
     c.bench_function(

arrow/src/util/bench_util.rs

@@ -491,6 +491,80 @@ where
     GenericListArray::<O>::from_iter_primitive::<T, _, _>(values)
 }
 
+/// Create a List/LargeList Array of primitive values using a fixed seed
+///
+/// See [`create_primitive_list_array_with_seed`] for details on arguments.
+pub fn create_primitive_list_array<O, T>(
+    size: usize,
+    null_density: f32,
+    list_null_density: f32,
+    max_list_size: usize,
+) -> GenericListArray<O>
+where
+    O: OffsetSizeTrait,
+    T: ArrowPrimitiveType,
+    StandardUniform: Distribution<T::Native>,
+{
+    let mut rng = seedable_rng();
+
+    let values = (0..size).map(|_| {
+        if rng.random::<f32>() < null_density {
+            None
+        } else {
+            let list_size = rng.random_range(0..=max_list_size);
+            let list_values: Vec<Option<T::Native>> = (0..list_size)
+                .map(|_| {
+                    if rng.random::<f32>() < list_null_density {
+                        None
+                    } else {
+                        Some(rng.random())
+                    }
+                })
+                .collect();
+            Some(list_values)
+        }
+    });
+
+    GenericListArray::<O>::from_iter_primitive::<T, _, _>(values)
+}
+
+/// Create a ListViewArray of primitive values using a fixed seed
+///
+/// See [`create_primitive_list_array_with_seed`] for details on arguments.
+pub fn create_primitive_list_view_array<O, T>(
+    size: usize,
+    null_density: f32,
+    list_null_density: f32,
+    max_list_size: usize,
+) -> GenericListViewArray<O>
+where
+    T: ArrowPrimitiveType,
+    StandardUniform: Distribution<T::Native>,
+    O: OffsetSizeTrait,
+{
+    let mut rng = seedable_rng();
+
+    let values = (0..size).map(|_| {
+        if rng.random::<f32>() < null_density {
+            None
+        } else {
+            let list_size = rng.random_range(0..=max_list_size);
+            let list_values: Vec<Option<T::Native>> = (0..list_size)
+                .map(|_| {
+                    if rng.random::<f32>() < list_null_density {
+                        None
+                    } else {
+                        Some(rng.random())
+                    }
+                })
+                .collect();
+            Some(list_values)
+        }
+    });
+
+    GenericListViewArray::<O>::from_iter_primitive::<T, _, _>(values)
+}
+
 /// Create primitive run array for given logical and physical array lengths
 pub fn create_primitive_run_array<R: RunEndIndexType, V: ArrowPrimitiveType>(
     logical_array_len: usize,