Skip to content

refactor: use arrow make_comparator for nested structural equality in arrays_overlap and array_position #5176

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andygrove merged 1 commit into from
Aug 1, 2026
+106 −110
Merged

refactor: use arrow make_comparator for nested structural equality in arrays_overlap and array_position #5176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 47 additions & 10 deletions native/spark-expr/src/array_funcs/array_position.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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),
}
}
Expand Down Expand Up @@ -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>,
Expand All @@ -270,28 +273,62 @@ 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)) {
continue;
}
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;
}
}
}

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,
Expand Down
159 changes: 59 additions & 100 deletions native/spark-expr/src/array_funcs/arrays_overlap.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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(_)
Expand All @@ -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;

Expand Down Expand Up @@ -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
Expand Down
Loading