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,7 +15,7 @@
 // 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};
@@ -71,6 +71,41 @@ 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),
+    )
+}
+
+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 +148,31 @@ 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))))
+    });
+
+    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);

native/spark-expr/src/array_funcs/arrays_overlap.rs

@@ -223,6 +223,29 @@ fn arrays_overlap_list<OffsetSize: OffsetSizeTrait>(
             left_values.as_string::<i64>(),
             right_values.as_string::<i64>()
         ),
+        dt if needs_comparator(dt) => {
+            // Spark's nested path compares with ordering.equiv, where -0.0 == 0.0 and every NaN
+            // is equal, but Arrow's comparator uses total order. Normalize float leaves once per
+            // column so the comparator built over the full child arrays matches Spark.
+            let (left_values, right_values) = if has_float_leaf(dt) {
+                (
+                    normalize_nested_floats(left_values),
+                    normalize_nested_floats(right_values),
+                )
+            } else {
+                (Arc::clone(left_values), Arc::clone(right_values))
+            };
+            let comparator = make_comparator(
+                left_values.as_ref(),
+                right_values.as_ref(),
+                SortOptions::default(),
+            )?;
+            Ok(overlap_rows(
+                left,
+                right,
+                nested_row_overlap(&left_values, &right_values, comparator.as_ref()),
+            ))
+        }
         _ => arrays_overlap_list_generic(left, right),
     }
 }
@@ -397,6 +420,30 @@ where
     }
 }
 
+/// Row overlap for nested element types using one comparator for the full child arrays.
+fn nested_row_overlap<'a>(
+    left: &'a ArrayRef,
+    right: &'a ArrayRef,
+    comparator: &'a dyn Fn(usize, usize) -> Ordering,
+) -> impl FnMut(Range<usize>, Range<usize>) -> bool + 'a {
+    move |left_range, right_range| {
+        for li in left_range {
+            if left.is_null(li) {
+                continue;
+            }
+            for ri in right_range.clone() {
+                if right.is_null(ri) {
+                    continue;
+                }
+                if comparator(li, ri) == Ordering::Equal {
+                    return true;
+                }
+            }
+        }
+        false
+    }
+}
+
 fn normalize_list_element_floats<OffsetSize: OffsetSizeTrait>(
     list: &GenericListArray<OffsetSize>,
 ) -> GenericListArray<OffsetSize> {
@@ -413,10 +460,11 @@ fn normalize_list_element_floats<OffsetSize: OffsetSizeTrait>(
     )
 }
 
-/// Fallback for nested and otherwise unhandled element types.
+/// Fallback for otherwise unhandled element types, including nested elements whose two sides
+/// have different data types.
 ///
 /// note: Spark's flat arrays_overlap (HashSet<Double>) treats -0.0 and 0.0 as different,
-/// only the nested path here treats them as equal. this normalization can't move into the
+/// only the nested path treats them as equal. this normalization can't move into the
 /// flat fast path in arrays_overlap_list without breaking that difference.
 fn arrays_overlap_list_generic<OffsetSize: OffsetSizeTrait>(
     left: &GenericListArray<OffsetSize>,
@@ -464,26 +512,12 @@ fn arrays_overlap_list_generic<OffsetSize: OffsetSizeTrait>(
             (&right_values, &left_values)
         };
 
-        let comparator = if needs_comparator(probe.data_type()) {
-            Some(make_comparator(
-                probe.as_ref(),
-                search.as_ref(),
-                SortOptions::default(),
-            )?)
-        } else {
-            None
-        };
-
         for pi in 0..probe.len() {
             if probe.is_null(pi) {
                 has_null = true;
                 continue;
             }
-            let (found, null_eq) = if let Some(comparator) = &comparator {
-                find_in_array_nested(pi, search, comparator.as_ref())
-            } else {
-                find_in_array_flat(probe, pi, search)?
-            };
+            let (found, null_eq) = find_in_array_flat(probe, pi, search)?;
             if null_eq {
                 has_null = true;
             }
@@ -513,25 +547,6 @@ fn find_in_array_flat(probe: &ArrayRef, pi: usize, search: &ArrayRef) -> Result<
     Ok((eq_result.true_count() > 0, eq_result.null_count() > 0))
 }
 
-/// Element-by-element search using Arrow's nested comparator.
-fn find_in_array_nested(
-    pi: usize,
-    search: &ArrayRef,
-    comparator: &dyn Fn(usize, usize) -> Ordering,
-) -> (bool, bool) {
-    let mut has_null = false;
-    for si in 0..search.len() {
-        if search.is_null(si) {
-            has_null = true;
-            continue;
-        }
-        if comparator(pi, si) == Ordering::Equal {
-            return (true, has_null);
-        }
-    }
-    (false, has_null)
-}
-
 fn needs_comparator(dt: &DataType) -> bool {
     matches!(
         dt,
@@ -895,6 +910,47 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_nested_array_sliced_offsets_and_nulls() -> Result<()> {
+        let make_rows = |rows: &[&[Option<&[i32]>]]| {
+            let mut builder = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
+            for row in rows {
+                for element in *row {
+                    if let Some(values) = element {
+                        builder.values().values().append_slice(values);
+                        builder.values().append(true);
+                    } else {
+                        builder.values().append(false);
+                    }
+                }
+                builder.append(true);
+            }
+            builder.finish()
+        };
+        let left = make_rows(&[
+            &[Some(&[999])],
+            &[Some(&[10])],
+            &[Some(&[10]), None],
+            &[Some(&[50]), Some(&[60]), Some(&[70])],
+        ])
+        .slice(1, 3);
+        let right = make_rows(&[
+            &[Some(&[999])],
+            &[Some(&[20]), Some(&[30]), Some(&[40])],
+            &[Some(&[20])],
+            &[Some(&[60])],
+        ])
+        .slice(1, 3);
+
+        let result = arrays_overlap_list::<i32>(&left, &right)?;
+        let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+        assert_eq!(
+            result,
+            &BooleanArray::from(vec![Some(false), None, Some(true)])
+        );
+        Ok(())
+    }
+
     #[test]
     fn test_nested_array_basic_overlap() -> Result<()> {
         // [[1,2], [3,4]] vs [[3,4], [5,6]] => true