diff --git a/native/spark-expr/benches/arrays_overlap.rs b/native/spark-expr/benches/arrays_overlap.rs index c20efdd8b6..67f5b56d13 100644 --- a/native/spark-expr/benches/arrays_overlap.rs +++ b/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); diff --git a/native/spark-expr/src/array_funcs/arrays_overlap.rs b/native/spark-expr/src/array_funcs/arrays_overlap.rs index 542687c7d1..e23a010b63 100644 --- a/native/spark-expr/src/array_funcs/arrays_overlap.rs +++ b/native/spark-expr/src/array_funcs/arrays_overlap.rs @@ -223,6 +223,29 @@ fn arrays_overlap_list( left_values.as_string::(), right_values.as_string::() ), + 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, Range) -> 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( list: &GenericListArray, ) -> GenericListArray { @@ -413,10 +460,11 @@ fn normalize_list_element_floats( ) } -/// 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) 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( left: &GenericListArray, @@ -464,26 +512,12 @@ fn arrays_overlap_list_generic( (&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::(&left, &right)?; + let result = result.as_any().downcast_ref::().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