This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new ba21f02dec perf: use Arrow comparator for nested array equality (#5176)
ba21f02dec is described below
commit ba21f02dec7104d8c10b6e925b3729d6600c4522
Author: Peter Lee <[email protected]>
AuthorDate: Sat Aug 1 22:43:01 2026 +0800
perf: use Arrow comparator for nested array equality (#5176)
---
.../spark-expr/src/array_funcs/array_position.rs | 57 ++++++--
.../spark-expr/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<O: OffsetSizeTrait>(
}
DataType::Utf8 => position_string::<O, i32>(list_array, offsets,
values, element),
DataType::LargeUtf8 => position_string::<O, i64>(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::<O>(list_array, offsets, values, element),
}
}
@@ -260,7 +263,7 @@ fn position_string<O: OffsetSizeTrait, S: OffsetSizeTrait>(
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<O: OffsetSizeTrait>(
list_array: &GenericListArray<O>,
offsets: &arrow::buffer::OffsetBuffer<O>,
@@ -270,6 +273,7 @@ fn position_fallback<O: OffsetSizeTrait>(
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<O: OffsetSizeTrait>(
}
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<O: OffsetSizeTrait>(
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::<Float64Type, _, _>([
+ 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::<Float64Type, _, _>([
+ 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::<Int64Array>().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<OffsetSize:
OffsetSizeTrait>(
(&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<bool> {
- // 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::<GenericListArray<i32>>()
- .unwrap();
- let rl = right
- .as_any()
- .downcast_ref::<GenericListArray<i32>>()
- .unwrap();
- list_structural_eq(&ll.value(li), &rl.value(ri))
- }
- DataType::LargeList(_) => {
- let ll = left
- .as_any()
- .downcast_ref::<GenericListArray<i64>>()
- .unwrap();
- let rl = right
- .as_any()
- .downcast_ref::<GenericListArray<i64>>()
- .unwrap();
- list_structural_eq(&ll.value(li), &rl.value(ri))
- }
- DataType::FixedSizeList(_, _) => {
- let ll =
left.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
- let rl =
right.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
- list_structural_eq(&ll.value(li), &rl.value(ri))
- }
- DataType::Struct(_) => {
- let ls = left.as_any().downcast_ref::<StructArray>().unwrap();
- let rs = right.as_any().downcast_ref::<StructArray>().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<bool> {
- 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<bool> {
- 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::<i32>(&left, &right)?;
+ let result = result.as_any().downcast_ref::<BooleanArray>().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::<i32>(&left, &right)?;
+ let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+ assert!(!result.value(0));
+ Ok(())
+ }
+
#[test]
fn test_nested_array_basic_overlap() -> Result<()> {
// [[1,2], [3,4]] vs [[3,4], [5,6]] => true
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
