From 1e54383071f2efe2bbbfe21e8e6f975af65f57eb Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 1 Aug 2026 00:11:48 +0800 Subject: [PATCH] perf: use Arrow comparator for nested array equality --- .../src/array_funcs/array_position.rs | 57 +++++-- .../src/array_funcs/arrays_overlap.rs | 159 +++++++----------- 2 files changed, 106 insertions(+), 110 deletions(-) diff --git a/native/spark-expr/src/array_funcs/array_position.rs b/native/spark-expr/src/array_funcs/array_position.rs index a5841b27d3..191091aabf 100644 --- a/native/spark-expr/src/array_funcs/array_position.rs +++ b/native/spark-expr/src/array_funcs/array_position.rs @@ -16,9 +16,11 @@ // under the License. use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, GenericListArray, Int64Array, OffsetSizeTrait, + make_comparator, Array, ArrayRef, AsArray, BooleanArray, GenericListArray, Int64Array, + OffsetSizeTrait, }; use arrow::buffer::{NullBuffer, ScalarBuffer}; +use arrow::compute::SortOptions; use arrow::datatypes::{ ArrowPrimitiveType, DataType, Date32Type, Decimal128Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, TimestampMicrosecondType, @@ -28,6 +30,7 @@ use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; use num::Float; +use std::cmp::Ordering; use std::sync::Arc; /// Spark array_position() function that returns the 1-based position of an element in an array. @@ -102,7 +105,7 @@ fn generic_array_position( } DataType::Utf8 => position_string::(list_array, offsets, values, element), DataType::LargeUtf8 => position_string::(list_array, offsets, values, element), - // Fallback to ScalarValue for complex types (nested arrays, etc.) + // Fallback to Arrow's comparator for complex types (nested arrays, etc.) _ => position_fallback::(list_array, offsets, values, element), } } @@ -260,7 +263,7 @@ fn position_string( Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) } -/// Fallback for complex types (nested arrays, structs, etc.) using ScalarValue comparison. +/// Fallback for complex types (nested arrays, structs, etc.) using Arrow's comparator. fn position_fallback( list_array: &GenericListArray, offsets: &arrow::buffer::OffsetBuffer, @@ -270,6 +273,7 @@ fn position_fallback( let num_rows = list_array.len(); let nulls = combined_nulls(list_array.nulls(), element.nulls()); let mut result = vec![0i64; num_rows]; + let comparator = make_comparator(values.as_ref(), element.as_ref(), SortOptions::default())?; for (row_index, w) in offsets.windows(2).enumerate() { if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { @@ -277,14 +281,10 @@ fn position_fallback( } let start = w[0].as_usize(); let end = w[1].as_usize(); - let search_scalar = ScalarValue::try_from_array(element, row_index)?; for i in start..end { - if !values.is_null(i) { - let item_scalar = ScalarValue::try_from_array(values, i)?; - if search_scalar == item_scalar { - result[row_index] = (i - start + 1) as i64; - break; - } + if !values.is_null(i) && comparator(i, row_index) == Ordering::Equal { + result[row_index] = (i - start + 1) as i64; + break; } } } @@ -292,6 +292,43 @@ fn position_fallback( Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) } +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::ListArray; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Float64Type}; + + #[test] + fn test_nested_float_and_null_position() -> DataFusionResult<()> { + // Arrow and the previous ScalarValue fallback distinguish signed zeros, so the second + // row matches at position 2 rather than position 1. + let values = ListArray::from_iter_primitive::([ + Some(vec![Some(1.0)]), + Some(vec![Some(f64::NAN)]), + Some(vec![Some(-0.0)]), + Some(vec![Some(0.0)]), + Some(vec![Some(1.0), None]), + ]); + let array = ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 2, 4, 5].into()), + Arc::new(values), + None, + ); + let element = ListArray::from_iter_primitive::([ + Some(vec![Some(f64::NAN)]), + Some(vec![Some(0.0)]), + Some(vec![Some(1.0), None]), + ]); + + let result = array_position_inner(&[Arc::new(array), Arc::new(element)])?; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result, &Int64Array::from(vec![2, 2, 1])); + Ok(()) + } +} + #[derive(Debug, Hash, Eq, PartialEq)] pub struct SparkArrayPositionFunc { signature: Signature, diff --git a/native/spark-expr/src/array_funcs/arrays_overlap.rs b/native/spark-expr/src/array_funcs/arrays_overlap.rs index cd609f9680..bd75a6ddcc 100644 --- a/native/spark-expr/src/array_funcs/arrays_overlap.rs +++ b/native/spark-expr/src/array_funcs/arrays_overlap.rs @@ -30,11 +30,12 @@ //! - false if no overlap and neither array contains null elements use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, FixedSizeListArray, GenericListArray, - GenericStringArray, OffsetSizeTrait, PrimitiveArray, Scalar, StructArray, + make_comparator, Array, ArrayRef, AsArray, BooleanArray, GenericListArray, GenericStringArray, + OffsetSizeTrait, PrimitiveArray, Scalar, }; use arrow::buffer::NullBuffer; use arrow::compute::kernels::cmp::eq; +use arrow::compute::SortOptions; use arrow::datatypes::{ ArrowPrimitiveType, DataType, Date32Type, Date64Type, Decimal128Type, FieldRef, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, TimeUnit, TimestampMicrosecondType, @@ -45,6 +46,7 @@ use datafusion::common::{exec_err, utils::take_function_args, HashSet, Result, S use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; +use std::cmp::Ordering; use std::hash::Hash; use std::ops::Range; use std::sync::Arc; @@ -426,18 +428,25 @@ fn arrays_overlap_list_generic( (&right_values, &left_values) }; - // Check element type once outside the loop. - let use_vectorized = !needs_recursive_eq(probe.data_type()); + 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 use_vectorized { - find_in_array_flat(probe, pi, search)? + let (found, null_eq) = if let Some(comparator) = &comparator { + find_in_array_nested(pi, search, comparator.as_ref()) } else { - find_in_array_nested(probe, pi, search)? + find_in_array_flat(probe, pi, search)? }; if null_eq { has_null = true; @@ -468,22 +477,26 @@ 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 structural equality for nested types. -fn find_in_array_nested(probe: &ArrayRef, pi: usize, search: &ArrayRef) -> Result<(bool, bool)> { +/// 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 structural_eq(probe.as_ref(), pi, search.as_ref(), si)? { - return Ok((true, has_null)); + if comparator(pi, si) == Ordering::Equal { + return (true, has_null); } } - Ok((false, has_null)) + (false, has_null) } -fn needs_recursive_eq(dt: &DataType) -> bool { +fn needs_comparator(dt: &DataType) -> bool { matches!( dt, DataType::List(_) @@ -493,96 +506,12 @@ fn needs_recursive_eq(dt: &DataType) -> bool { ) } -/// Structural equality for array elements (grouping semantics: NULL == NULL is true). -/// This matches Spark's `ordering.equiv` used inside `arrays_overlap`. -/// Three-valued null logic only applies to outer-level null elements (handled by the caller). -fn structural_eq(left: &dyn Array, li: usize, right: &dyn Array, ri: usize) -> Result { - // NullArray::is_null() returns false (no null buffer), so check data type first. - if left.data_type() == &DataType::Null && right.data_type() == &DataType::Null { - return Ok(true); - } - - if left.is_null(li) && right.is_null(ri) { - return Ok(true); - } - if left.is_null(li) || right.is_null(ri) { - return Ok(false); - } - - match left.data_type() { - DataType::List(_) => { - let ll = left - .as_any() - .downcast_ref::>() - .unwrap(); - let rl = right - .as_any() - .downcast_ref::>() - .unwrap(); - list_structural_eq(&ll.value(li), &rl.value(ri)) - } - DataType::LargeList(_) => { - let ll = left - .as_any() - .downcast_ref::>() - .unwrap(); - let rl = right - .as_any() - .downcast_ref::>() - .unwrap(); - list_structural_eq(&ll.value(li), &rl.value(ri)) - } - DataType::FixedSizeList(_, _) => { - let ll = left.as_any().downcast_ref::().unwrap(); - let rl = right.as_any().downcast_ref::().unwrap(); - list_structural_eq(&ll.value(li), &rl.value(ri)) - } - DataType::Struct(_) => { - let ls = left.as_any().downcast_ref::().unwrap(); - let rs = right.as_any().downcast_ref::().unwrap(); - struct_structural_eq(ls, li, rs, ri) - } - _ => { - // Both non-null at this point; eq on two non-null scalars is definitive. - let l = Scalar::new(left.slice(li, 1)); - let r = Scalar::new(right.slice(ri, 1)); - let result = eq(&l, &r) - .map_err(|e| datafusion::error::DataFusionError::ArrowError(Box::new(e), None))?; - Ok(result.value(0)) - } - } -} - -fn list_structural_eq(left: &ArrayRef, right: &ArrayRef) -> Result { - if left.len() != right.len() { - return Ok(false); - } - for k in 0..left.len() { - if !structural_eq(left.as_ref(), k, right.as_ref(), k)? { - return Ok(false); - } - } - Ok(true) -} - -fn struct_structural_eq( - left: &StructArray, - li: usize, - right: &StructArray, - ri: usize, -) -> Result { - for (lc, rc) in left.columns().iter().zip(right.columns().iter()) { - if !structural_eq(lc.as_ref(), li, rc.as_ref(), ri)? { - return Ok(false); - } - } - Ok(true) -} - #[cfg(test)] mod tests { use super::*; - use arrow::array::{Int32Array, Int32Builder, ListArray, ListBuilder, StructBuilder}; + use arrow::array::{ + Float64Builder, Int32Array, Int32Builder, ListArray, ListBuilder, StructBuilder, + }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::Field; @@ -763,6 +692,36 @@ mod tests { outer_builder.finish() } + fn make_nested_float_list(elements: &[&[f64]]) -> ListArray { + let mut builder = ListBuilder::new(ListBuilder::new(Float64Builder::new())); + for element in elements { + for value in *element { + builder.values().values().append_value(*value); + } + builder.values().append(true); + } + builder.append(true); + builder.finish() + } + + #[test] + fn test_nested_float_total_order() -> Result<()> { + // Preserve the existing Arrow total-order behavior: NaN matches itself, while signed + // zeros are distinct. + let left = make_nested_float_list(&[&[f64::NAN]]); + let right = make_nested_float_list(&[&[f64::NAN]]); + let result = arrays_overlap_list::(&left, &right)?; + let result = result.as_any().downcast_ref::().unwrap(); + assert!(result.value(0)); + + let left = make_nested_float_list(&[&[0.0]]); + let right = make_nested_float_list(&[&[-0.0]]); + let result = arrays_overlap_list::(&left, &right)?; + let result = result.as_any().downcast_ref::().unwrap(); + assert!(!result.value(0)); + Ok(()) + } + #[test] fn test_nested_array_basic_overlap() -> Result<()> { // [[1,2], [3,4]] vs [[3,4], [5,6]] => true