refactor: use arrow make_comparator for nested structural equality in arrays_overlap and array_position [2/2]

native/spark-expr/benches/arrays_overlap.rs

@@ -15,14 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{ArrayRef, Int32Array, ListArray, StringArray};
+use arrow::array::{ArrayRef, Int32Array, ListArray, StringArray, StructArray};
 use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::datatypes::{DataType, Field};
 use criterion::{criterion_group, criterion_main, Criterion};
 use datafusion::common::config::ConfigOptions;
 use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
 use datafusion_comet_spark_expr::SparkArraysOverlap;
 use std::hint::black_box;
+use std::ops::Range;
 use std::sync::Arc;
 
 fn list_of(values: ArrayRef, rows: usize, elems_per_row: usize) -> ArrayRef {
@@ -71,6 +72,65 @@ fn string_lists(rows: usize, elems_per_row: usize, offset: usize) -> (ArrayRef,
     )
 }
 
+fn nested_int_lists(rows: usize, elems_per_row: usize, offset: i32) -> (ArrayRef, ArrayRef) {
+    let total = rows * elems_per_row;
+    let build = |value_offset: i32| {
+        let values: ArrayRef = Arc::new(Int32Array::from_iter_values(
+            (0..total).flat_map(|i| [0, 1, 2, i as i32 + value_offset]),
+        ));
+        list_of(values, total, 4)
+    };
+    (
+        list_of(build(0), rows, elems_per_row),
+        list_of(build(offset), rows, elems_per_row),
+    )
+}
+
+/// Nested int32 lists of one-value lists, none of them null. Each left row holds `[v]` for each
+/// `v` in `left`, and each right row likewise for `right`.
+fn nested_singleton_lists(
+    rows: usize,
+    left: Range<i32>,
+    right: Range<i32>,
+) -> (ArrayRef, ArrayRef) {
+    let build = |values: Range<i32>| {
+        let len = values.len();
+        let inner: ArrayRef = Arc::new(Int32Array::from_iter_values(
+            (0..rows).flat_map(|_| values.clone()),
+        ));
+        let offsets: Vec<i32> = (0..=rows * len).map(|i| i as i32).collect();
+        let singletons: ArrayRef = Arc::new(ListArray::new(
+            Arc::new(Field::new("item", DataType::Int32, true)),
+            OffsetBuffer::new(offsets.into()),
+            inner,
+            None,
+        ));
+        list_of(singletons, rows, len)
+    };
+    (build(left), build(right))
+}
+
+fn struct_lists(rows: usize, elems_per_row: usize) -> (ArrayRef, ArrayRef) {
+    let total = rows * elems_per_row;
+    let build = |offset: i32| -> ArrayRef {
+        let first: ArrayRef = Arc::new(Int32Array::from_value(0, total));
+        let second: ArrayRef = Arc::new(Int32Array::from_iter_values(
+            (0..total).map(|i| i as i32 + offset),
+        ));
+        Arc::new(StructArray::from(vec![
+            (Arc::new(Field::new("first", DataType::Int32, false)), first),
+            (
+                Arc::new(Field::new("second", DataType::Int32, false)),
+                second,
+            ),
+        ]))
+    };
+    (
+        list_of(build(0), rows, elems_per_row),
+        list_of(build(total as i32), rows, elems_per_row),
+    )
+}
+
 fn invoke(udf: &SparkArraysOverlap, left: &ArrayRef, right: &ArrayRef) -> ColumnarValue {
     udf.invoke_with_args(ScalarFunctionArgs {
         args: vec![
@@ -113,6 +173,45 @@ fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("spark_arrays_overlap: utf8 long lists", |b| {
         b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right))))
     });
+
+    let (left, right) = nested_int_lists(rows, 8, (rows * 8) as i32);
+    c.bench_function("spark_arrays_overlap: nested int32 short lists", |b| {
+        b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right))))
+    });
+
+    let (left, right) = nested_int_lists(64, 64, 64 * 64);
+    c.bench_function("spark_arrays_overlap: nested int32 long lists", |b| {
+        b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right))))
+    });
+
+    let (left, right) = nested_int_lists(rows, 8, 4);
+    c.bench_function("spark_arrays_overlap: nested int32 early match", |b| {
+        b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right))))
+    });
+
+    // Rows hold 128 and 64 elements with one match. The first bench puts it last in the longer
+    // side, the second last in the shorter side.
+    let (left, right) = nested_singleton_lists(256, 0..128, 127..191);
+    c.bench_function(
+        "spark_arrays_overlap: nested int32 unequal lengths, late match in longer side",
+        |b| b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right)))),
+    );
+
+    let (left, right) = nested_singleton_lists(256, 0..128, -63..1);
+    c.bench_function(
+        "spark_arrays_overlap: nested int32 unequal lengths, late match in shorter side",
+        |b| b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right)))),
+    );
+
+    let (left, right) = struct_lists(rows, 8);
+    c.bench_function("spark_arrays_overlap: nested struct short lists", |b| {
+        b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right))))
+    });
+
+    let (left, right) = struct_lists(64, 64);
+    c.bench_function("spark_arrays_overlap: nested struct long lists", |b| {
+        b.iter(|| black_box(invoke(&udf, black_box(&left), black_box(&right))))
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);