Skip to content
4 changes: 2 additions & 2 deletions docs/source/user-guide/latest/expressions.md
Original file line number Diff line number Diff line change
Expand Up @@ -158,12 +158,12 @@ The tables below list every Spark built-in expression with its current status.
| `array_join` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) |
| `array_max` | ✅ | Native | NaN ordering may differ ([details](compatibility/floating-point.md)) |
| `array_min` | ✅ | Native | NaN ordering may differ ([details](compatibility/floating-point.md)) |
| `array_position` | ✅ | Native | Binary/struct/map/null elements fall back |
| `array_position` | ✅ | Native | Binary/struct/map/null elements fall back; arrays with nested floating-point elements fall back to Spark by default, and the incompatible native path is opt-in via allowIncompatible ([#5191](https://github.com/apache/datafusion-comet/issues/5191)) |
| `array_prepend` | ✅ | — | |
| `array_remove` | ✅ | Native | |
| `array_repeat` | ✅ | Native | |
| `array_union` | ✅ | Native | NaN/signed-zero handling may differ ([details](compatibility/floating-point.md)) |
| `arrays_overlap` | ✅ | Native | |
| `arrays_overlap` | ✅ | Native | Arrays with nested floating-point elements fall back to Spark by default, and the incompatible native path is opt-in via allowIncompatible ([#5191](https://github.com/apache/datafusion-comet/issues/5191)) |
| `arrays_zip` | ✅ | Native | |
| `element_at` | ✅ | Native | |
| `flatten` | ✅ | Native | Binary/struct/map elements fall back |
Expand Down
62 changes: 61 additions & 1 deletion native/spark-expr/benches/arrays_overlap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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![
Expand Down Expand Up @@ -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);
Expand Down
5 changes: 3 additions & 2 deletions native/spark-expr/src/array_funcs/array_position.rs
Original file line number Diff line number Diff line change
Expand Up @@ -301,8 +301,9 @@ mod tests {

#[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.
// Signed-zero equality does not yet match Spark, so the second row matches at position 2;
// see https://github.com/apache/datafusion-comet/issues/5191.
// NaN and inner-null equality match Spark.
let values = ListArray::from_iter_primitive::<Float64Type, _, _>([
Some(vec![Some(1.0)]),
Some(vec![Some(f64::NAN)]),
Expand Down
118 changes: 81 additions & 37 deletions native/spark-expr/src/array_funcs/arrays_overlap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,18 @@ fn arrays_overlap_list<OffsetSize: OffsetSizeTrait>(
left_values.as_string::<i64>(),
right_values.as_string::<i64>()
),
dt if needs_comparator(dt) => {
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),
}
}
Expand Down Expand Up @@ -395,7 +407,31 @@ where
}
}

/// Fallback for nested and otherwise unhandled element types.
/// 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
}
}

/// Fallback for otherwise unhandled element types.
fn arrays_overlap_list_generic<OffsetSize: OffsetSizeTrait>(
left: &GenericListArray<OffsetSize>,
right: &GenericListArray<OffsetSize>,
Expand Down Expand Up @@ -435,26 +471,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;
}
Expand Down Expand Up @@ -484,25 +506,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,
Expand Down Expand Up @@ -839,14 +842,14 @@ mod tests {

#[test]
fn test_nested_float_total_order() -> Result<()> {
// Preserve the existing Arrow total-order behavior: NaN matches itself, while signed
// zeros are distinct.
// NaN equality matches Spark.
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));

// Signed-zero equality does not yet match Spark; see #5191.
let left = make_nested_float_list(&[&[0.0]]);
let right = make_nested_float_list(&[&[-0.0]]);
let result = arrays_overlap_list::<i32>(&left, &right)?;
Expand All @@ -855,6 +858,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
Expand Down
36 changes: 34 additions & 2 deletions spark/src/main/scala/org/apache/comet/serde/arrays.scala
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,18 @@ object CometArrayMin extends CometExpressionSerde[ArrayMin] {
}
}

object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] {
object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] with ArraysBase {

override def getIncompatibleReasons(): Seq[String] = Seq(nestedFloatIncompatibilityReason)

override def getSupportLevel(expr: ArraysOverlap): SupportLevel = {
if (hasNestedFloatElements(expr.left.dataType)) {
Incompatible(Some(nestedFloatIncompatibilityReason))
} else {
Compatible()
}
}

override def convert(
expr: ArraysOverlap,
inputs: Seq[Attribute],
Expand Down Expand Up @@ -730,12 +741,18 @@ object CometSize extends CometExpressionSerde[Size] {

object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with ArraysBase {

override def getIncompatibleReasons(): Seq[String] = Seq(nestedFloatIncompatibilityReason)

override def getSupportLevel(expr: ArrayPosition): SupportLevel = {
if (expr.children.forall(_.foldable)) {
// Fall back to Spark for all-literal args so ConstantFolding can handle it.
Unsupported(Some("all arguments are literals, falling back to Spark"))
} else {
childTypesSupportLevel(expr)
childTypesSupportLevel(expr) match {
case _: Compatible if hasNestedFloatElements(expr.left.dataType) =>
Incompatible(Some(nestedFloatIncompatibilityReason))
case level => level
}
}
}

Expand Down Expand Up @@ -825,6 +842,21 @@ object CometArraysZip extends CometExpressionSerde[ArraysZip] {

trait ArraysBase {

protected val nestedFloatIncompatibilityReason: String =
"Nested floating-point elements distinguish `-0.0` from `0.0`, unlike Spark " +
"(https://github.com/apache/datafusion-comet/issues/5191)"

protected def hasNestedFloatElements(dt: DataType): Boolean = dt match {
case ArrayType(elementType: ArrayType, _) =>
SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])
case ArrayType(elementType: StructType, _) =>
SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])
// Map elements cannot reach here: ArraysOverlap and ArrayPosition call
// TypeUtils.checkForOrderingExpr on the element type, and maps are not orderable,
// so Spark's analyzer rejects them before planning.
case _ => false
}

def isTypeSupported(dt: DataType): Boolean = {
import DataTypes._
dt match {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,16 @@ INSERT INTO test_overlap_nested VALUES (array(array(1, 2), array(3, 4)), array(a
query
SELECT a, b, arrays_overlap(a, b) FROM test_overlap_nested

-- nested floating-point signed-zero behavior differs from Spark
statement
CREATE TABLE test_overlap_nested_dbl(a array<array<double>>, b array<array<double>>) USING parquet

statement
INSERT INTO test_overlap_nested_dbl VALUES (array(array(0.0D)), array(array(-0.0D))), (array(array(double('NaN'))), array(array(double('NaN'))))

query ignore(https://github.com/apache/datafusion-comet/issues/5191)
SELECT a, b, arrays_overlap(a, b) FROM test_overlap_nested_dbl

-- struct element arrays
statement
CREATE TABLE test_overlap_struct(a array<struct<x:int, y:int>>, b array<struct<x:int, y:int>>) USING parquet
Expand Down
Loading