You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/02 09:32:42 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6518-1b488623a167d9c0ef325e2e48590ed7d81d8c46 created (now 4f459dbbbb)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6518-1b488623a167d9c0ef325e2e48590ed7d81d8c46
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 4f459dbbbb fix: match Spark's float semantics in array_remove and sort_array (#6518)
This branch includes the following new commits:
new 4f459dbbbb fix: match Spark's float semantics in array_remove and sort_array (#6518)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: match Spark's float semantics in array_remove and sort_array (#6518)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6518-1b488623a167d9c0ef325e2e48590ed7d81d8c46
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 4f459dbbbb744e77f829d0085b180e7201457e69
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri Oct 2 09:32:19 2026 +0000
fix: match Spark's float semantics in array_remove and sort_array (#6518)
* fix: match Spark's float semantics in array_remove and sort_array
Spark compares floats in array_remove with genEqual, where -0.0 equals
0.0 and all NaNs are equal, and sorts them in sort_array with a stable
sort in its SQL ordering, where NaN is the largest value. DataFusion's
array_remove_all and array_sort compare by IEEE 754 total order instead.
Replace array_remove_all in Comet's registry with SparkArrayRemove,
which handles elements holding a float at any depth and delegates other
types to DataFusion. Send float sort_array to a new spark_sort_array,
which takes Spark's containsNull so that an ascending sort of non-null
FLOAT or DOUBLE elements follows java.util.Arrays.sort, as Spark's
generated code does. Float sort_array no longer needs the codegen
dispatcher in strict floating-point mode.
Part of #6385.
* fix: return errors for mismatched types and cover more array shapes
Dispatch the float fast paths of array_remove on the element and value
types together, so a value of another type fails in spark_equality
instead of panicking, and return an error from sort_array for input
that is not a list. Add SQL cases for a literal array with a column
value and for null arrays, and benchmark array_remove with a value per
row.
* refactor: route array_remove in Scala and resolve the fast-sort rule there
Send float array_remove to a new spark_array_remove from the serde, as
sort_array already chooses its function there, which removes the wrapper
around DataFusion's array_remove_all. Resolve Spark's canPerformFastSort
rule in CometSortArray and pass whether -0.0 sorts first, and move that
ordering into float_semantics as compare_floats_java.
Count kept elements with one running popcount instead of a bit count per
row, which cuts array_remove of a constant from 77 to 26 us for 8192
rows of 8 doubles, and compact sort_array's valid values toward the end
of an ascending row so their count needs no separate pass. The nested
sort relies on spark_comparator's null ordering, and both functions
share their list and test helpers. The new fixtures no longer run twice
under a strict floating-point matrix that no query reads.
---
.../expression-audits/array_funcs.md | 10 +-
.../latest/compatibility/floating-point.md | 12 +
docs/source/user-guide/latest/expressions.md | 2 +-
docs/source/user-guide/latest/tuning/operators.md | 5 +-
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/float_arrays.rs | 165 +++++++
native/spark-expr/src/array_funcs/array_remove.rs | 483 +++++++++++++++++++
native/spark-expr/src/array_funcs/mod.rs | 72 +++
native/spark-expr/src/array_funcs/sort_array.rs | 521 +++++++++++++++++++++
native/spark-expr/src/comet_scalar_funcs.rs | 12 +-
native/spark-expr/src/float_semantics/mod.rs | 18 +-
.../main/scala/org/apache/comet/CometConf.scala | 8 +-
.../main/scala/org/apache/comet/serde/arrays.scala | 64 ++-
.../array/array_remove_floating_point.sql | 63 +++
.../expressions/array/routing_arrays_disabled.sql | 3 +-
.../expressions/array/routing_arrays_enabled.sql | 3 +-
.../array/sort_array_floating_point.sql | 68 +++
.../expressions/array/sort_array_strict_fp.sql | 7 +-
18 files changed, 1470 insertions(+), 50 deletions(-)
diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md
index cd167885d0..0911c45aff 100644
--- a/docs/source/contributor-guide/expression-audits/array_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/array_funcs.md
@@ -122,9 +122,11 @@
## array_remove
- Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8.
-- Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayRemove(left, right)`; removes all occurrences equal to `right`. Wired as `CometScalarFunction("array_remove")`. Falls back via `ArraysBase.isTypeSupported` for binary/struct/map/null child types.
+- Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayRemove(left, right)`; removes all occurrences equal to `right`. Falls back via `ArraysBase.isTypeSupported` for binary/struct/map/null child types.
- Spark 4.0.1 (audited 2026-05-27): `NullIntolerant` -> `nullIntolerant` field refactor.
- Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1.
+- Current status: `CometArrayRemove` sends arrays whose elements hold a `FLOAT` or `DOUBLE` at any depth to the native `spark_array_remove`, which compares as Spark's `genEqual` does (`-0.0` equals `0.0`, all NaNs are equal, nested elements through `spark_equality`) and keeps the bits of the elements it keeps. Other element types use DataFusion's `array_remove_all`. Null elements stay, and a null array or value gives null.
+- Performance (tuned 2026-10-01, PR [#6518](https://github.com/apache/datafusion-comet/pull/6518)): with a constant value, `FLOAT` and `DOUBLE` elements get one keep mask over all the values from `BooleanBuffer::collect_bool`, with the inverted validity ORed in, one running popcount over the mask's words for each row's kept count, and one `filter`. 41-76% less time than DataFusion's `array_remove_all` with a constant value, and 47-95% less with a value per row. Benchmark: `benches/float_ [...]
## array_repeat
@@ -196,8 +198,10 @@
## sort_array
- Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8.
-- Spark 3.5.8 (audited 2026-05-27): baseline. `SortArray(base, ascendingOrder) extends BinaryExpression with ArraySortLike`; the second arg must be a `Literal(_: Boolean, BooleanType)`. Comet `CometSortArray` flags `Incompatible` under strict floating-point and falls back for nested arrays whose innermost element is `Struct` or `Null`.
-- Spark 4.0.1 (audited 2026-05-27): trait set changes substantively: `ArraySortLike` and `NullIntolerant` are removed, `nullIntolerant = true` becomes an override, and `ascendingOrder` is widened to accept any foldable boolean (not just `Literal`). Comet's `CometSortArray` still requires a `Literal`, so the new foldable form falls back at convert time.
+- Spark 3.5.8 (audited 2026-05-27): baseline. `SortArray(base, ascendingOrder) extends BinaryExpression with ArraySortLike`; the second arg must be a `Literal(_: Boolean, BooleanType)`.
+- Spark 4.0.1 (audited 2026-05-27): trait set changes substantively: `ArraySortLike` and `NullIntolerant` are removed, `nullIntolerant = true` becomes an override, and `ascendingOrder` is widened to accept any foldable boolean (not just `Literal`).
- Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1.
+- Current status: elements with a `FLOAT` or `DOUBLE` at any depth go to the native `spark_sort_array`, which sorts as Spark's generated code does: a stable sort in Spark's SQL ordering, except that for an ascending sort of `FLOAT` or `DOUBLE` elements that cannot be null, the serde has it put `-0.0` before `0.0`, as `java.util.Arrays.sort` does. This holds in strict floating-point mode too. Other supported element types use DataFusion's `array_sort`, and unsupported ones route through t [...]
+- Performance (tuned 2026-10-01, PR [#6518](https://github.com/apache/datafusion-comet/pull/6518)): `FLOAT` and `DOUBLE` rows sort within one copy of all the values, each row's valid values copied without branching next to its nulls. Rows of up to 20 elements use the stable `sort_by`. Longer rows use `sort_unstable_by` and then put the zero and NaN runs back in their original order, because Rust's stable sort is up to 1.8 times slower between 33 and 63 elements. From 2% more to 21% less [...]
[Spark Expression Support]: ../../user-guide/latest/expressions.md
diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md
index 57a9ffe8a8..7ed704bb54 100644
--- a/docs/source/user-guide/latest/compatibility/floating-point.md
+++ b/docs/source/user-guide/latest/compatibility/floating-point.md
@@ -69,3 +69,15 @@ greater than non-NaN values. The original first equal element is retained: for e
`array_min(array(0.0D, -0.0D))` returns `0.0`, while reversing those elements returns `-0.0`.
The same ordering applies recursively to floating-point fields in arrays and structs. These
expressions do not require Spark's codegen dispatcher for floating-point compatibility.
+
+## `array_remove` and `sort_array`
+
+`array_remove` compares `FLOAT` and `DOUBLE` elements as Spark does: `-0.0` equals `0.0`, and all
+NaN representations are equal, inside nested arrays too. The elements it keeps retain their
+original NaN representations and zero signs.
+
+`sort_array` sorts in Spark's order, in which NaN sorts above every other value and all NaN
+representations tie, and keeps equal elements in their original order. In an array whose elements
+can be null, such as one built from nullable columns, `-0.0` and `0.0` therefore tie. Sorting an
+array whose elements cannot be null in ascending order, Spark's generated code puts `-0.0` before
+`0.0`, and so does Comet. Both expressions run natively in strict floating-point mode.
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 0cb4d96cf6..68372aeb17 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -176,7 +176,7 @@ The tables below list every Spark built-in expression with its current status.
| `sequence` | ✅ | Hybrid | Integral types run natively; date/timestamp sequences use codegen dispatch |
| `shuffle` | ✅ | Native | Binary/struct/map elements fall back |
| `slice` | ✅ | Native | Native ([#4149](https://github.com/apache/datafusion-comet/pull/4149)) |
-| `sort_array` | ✅ | Hybrid | Struct, nested-array, and null elements run natively; other element types (for example intervals), and floating-point elements when `spark.comet.exec.strictFloatingPoint=true`, route through the JVM codegen dispatcher |
+| `sort_array` | ✅ | Hybrid | Struct, nested-array, floating-point, and null elements run natively; other element types (for example intervals) route through the JVM codegen dispatcher |
---
diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md
index 38fb80f22e..f36546c237 100644
--- a/docs/source/user-guide/latest/tuning/operators.md
+++ b/docs/source/user-guide/latest/tuning/operators.md
@@ -160,6 +160,5 @@ an edge case that is not of concern for many users. Setting `spark.comet.exec.st
nested cases fall back to Spark, and they can be forced back onto the native path with
`spark.comet.expression.SortOrder.allowIncompatible=true`.
-`sort_array` is separate. It sorts array elements rather than ordering rows, and its elements are compared with Arrow's
-raw total ordering, so `spark.comet.exec.strictFloatingPoint=true` makes it fall back even for a scalar floating-point
-element type. Use `spark.comet.expression.SortArray.allowIncompatible=true` to keep it native.
+`sort_array` sorts array elements rather than ordering rows. It follows Spark's floating-point ordering as well, so it
+also stays native with `spark.comet.exec.strictFloatingPoint=true`.
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 38f10208bb..a97334b56f 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -453,3 +453,7 @@ harness = false
[[bench]]
name = "nested_comparison"
harness = false
+
+[[bench]]
+name = "float_arrays"
+harness = false
diff --git a/native/spark-expr/benches/float_arrays.rs b/native/spark-expr/benches/float_arrays.rs
new file mode 100644
index 0000000000..44a4d78a2f
--- /dev/null
+++ b/native/spark-expr/benches/float_arrays.rs
@@ -0,0 +1,165 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Compare Comet's float `sort_array` and `array_remove` with DataFusion's `array_sort` and
+//! `array_remove_all` on doubles without zeros or NaNs, where both return the same arrays, which
+//! is checked before timing; special-value semantics belong in tests.
+
+use arrow::array::{Array, ArrayRef, Float64Array, ListArray};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::ScalarValue;
+use datafusion::functions_nested::remove::array_remove_all_udf;
+use datafusion::functions_nested::sort::array_sort_udf;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF};
+use datafusion_comet_spark_expr::{SparkArrayRemove, SparkSortArray};
+use std::hint::black_box;
+use std::sync::Arc;
+use std::time::Duration;
+
+const ROWS: usize = 8192;
+
+/// The value of element `i`. Values repeat so that `array_remove` finds something to remove.
+fn value(i: usize) -> f64 {
+ ((i * 7919) % 101) as f64 * 0.5 + 1.0
+}
+
+/// `ROWS` lists of `len` doubles, with every tenth list and every seventh element null if
+/// `nulls`.
+fn lists(len: usize, nulls: bool) -> ArrayRef {
+ let values = if nulls {
+ Float64Array::from_iter((0..ROWS * len).map(|i| (i % 7 != 0).then(|| value(i))))
+ } else {
+ Float64Array::from_iter_values((0..ROWS * len).map(value))
+ };
+ Arc::new(ListArray::new(
+ Arc::new(Field::new_list_field(DataType::Float64, true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(len, ROWS)),
+ Arc::new(values),
+ nulls.then(|| NullBuffer::from_iter((0..ROWS).map(|row| row % 10 != 0))),
+ ))
+}
+
+fn args(args: Vec<ColumnarValue>, list: &ArrayRef) -> ScalarFunctionArgs {
+ ScalarFunctionArgs {
+ arg_fields: args
+ .iter()
+ .map(|arg| Arc::new(Field::new("arg", arg.data_type(), true)))
+ .collect(),
+ args,
+ number_rows: ROWS,
+ return_field: Arc::new(Field::new("result", list.data_type().clone(), true)),
+ config_options: Arc::new(ConfigOptions::default()),
+ }
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("float_arrays");
+ group.sample_size(20);
+ group.warm_up_time(Duration::from_millis(250));
+ group.measurement_time(Duration::from_secs(1));
+ let comet_sort = ScalarUDF::new_from_impl(SparkSortArray::default());
+ let comet_remove = ScalarUDF::new_from_impl(SparkArrayRemove::default());
+ let (datafusion_sort, datafusion_remove) = (array_sort_udf(), array_remove_all_udf());
+ for len in [8, 50] {
+ for nulls in [false, true] {
+ let list = lists(len, nulls);
+ let data = format!("len={len}_null={nulls}");
+ let column = ColumnarValue::Array(Arc::clone(&list));
+ let scalar = |value: ScalarValue| ColumnarValue::Scalar(value);
+ let comet_sort_args = args(
+ vec![
+ column.clone(),
+ scalar(ScalarValue::Boolean(Some(true))),
+ scalar(ScalarValue::Boolean(Some(false))),
+ ],
+ &list,
+ );
+ let datafusion_sort_args = args(
+ vec![
+ column.clone(),
+ scalar(ScalarValue::from("ASC")),
+ scalar(ScalarValue::from("NULLS FIRST")),
+ ],
+ &list,
+ );
+ let remove_args = args(
+ vec![column.clone(), scalar(ScalarValue::Float64(Some(26.0)))],
+ &list,
+ );
+ // A value per row: the row's second element.
+ let row_values =
+ Float64Array::from_iter_values((0..ROWS).map(|row| value(row * len + 1)));
+ let remove_column_args = args(
+ vec![column, ColumnarValue::Array(Arc::new(row_values))],
+ &list,
+ );
+ let cases = [
+ (
+ "sort_array",
+ &comet_sort,
+ &datafusion_sort,
+ &comet_sort_args,
+ &datafusion_sort_args,
+ ),
+ (
+ "array_remove",
+ &comet_remove,
+ &datafusion_remove,
+ &remove_args,
+ &remove_args,
+ ),
+ (
+ "array_remove_column",
+ &comet_remove,
+ &datafusion_remove,
+ &remove_column_args,
+ &remove_column_args,
+ ),
+ ];
+ for (op, comet, datafusion, comet_args, datafusion_args) in cases {
+ let evaluate = |udf: &ScalarUDF, args: &ScalarFunctionArgs| {
+ udf.invoke_with_args(args.clone())
+ .unwrap()
+ .into_array(ROWS)
+ .unwrap()
+ };
+ assert_eq!(
+ evaluate(comet, comet_args).to_data(),
+ evaluate(datafusion, datafusion_args).to_data(),
+ "{op} {data}"
+ );
+ for (engine, udf, udf_args) in [
+ ("comet", comet, comet_args),
+ ("datafusion", datafusion.as_ref(), datafusion_args),
+ ] {
+ group.bench_function(BenchmarkId::new(format!("{op}_{engine}"), &data), |b| {
+ b.iter(|| {
+ black_box(udf.invoke_with_args(black_box(udf_args.clone())).unwrap())
+ })
+ });
+ }
+ }
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/src/array_funcs/array_remove.rs b/native/spark-expr/src/array_funcs/array_remove.rs
new file mode 100644
index 0000000000..aac6dd4c8c
--- /dev/null
+++ b/native/spark-expr/src/array_funcs/array_remove.rs
@@ -0,0 +1,483 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use super::with_values;
+use crate::float_semantics::{compare_floats, spark_equality};
+use arrow::array::{
+ Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, ListArray, PrimitiveArray,
+};
+use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
+use arrow::compute::filter;
+use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type};
+use datafusion::common::{exec_err, utils::take_function_args, Result, ScalarValue};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use num::Float;
+
+/// Spark's `array_remove` for arrays whose elements hold floats at any depth.
+///
+/// Spark removes every element equal to the value under `genEqual`, in which `-0.0` equals `0.0`
+/// and all NaNs are equal, at any depth of an array or struct element. It keeps null elements,
+/// and returns null when the array or the value is null. DataFusion's `array_remove_all`, which
+/// Comet uses for other element types, compares the bits instead, so it keeps a `-0.0` when
+/// removing `0.0`, and a NaN whose bits differ.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SparkArrayRemove {
+ signature: Signature,
+}
+
+impl Default for SparkArrayRemove {
+ fn default() -> Self {
+ Self {
+ signature: Signature::any(2, Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkArrayRemove {
+ fn name(&self) -> &str {
+ "spark_array_remove"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ Ok(arg_types[0].clone())
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ let [array, value] = take_function_args(self.name(), &args.args)?;
+ let all_scalars = matches!(
+ (array, value),
+ (ColumnarValue::Scalar(_), ColumnarValue::Scalar(_))
+ );
+ let rows = if all_scalars { 1 } else { args.number_rows };
+ let array = array.to_array(rows)?;
+ // Spark arrays use Arrow's 32-bit List layout.
+ let Some(list) = array.as_list_opt::<i32>() else {
+ return exec_err!("spark_array_remove takes a list, got {}", array.data_type());
+ };
+ let result = match value {
+ ColumnarValue::Scalar(needle) => remove_constant(list, needle)?,
+ ColumnarValue::Array(value) => array_remove(list, value)?,
+ };
+ if all_scalars {
+ Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &result, 0,
+ )?))
+ } else {
+ Ok(ColumnarValue::Array(result))
+ }
+ }
+}
+
+fn array_remove(array: &ListArray, value: &ArrayRef) -> Result<ArrayRef> {
+ // Spark casts the value to the element type. Anything else fails in `spark_equality`.
+ match (array.value_type(), value.data_type()) {
+ (DataType::Float32, DataType::Float32) => Ok(remove_floats::<Float32Type>(array, value)),
+ (DataType::Float64, DataType::Float64) => Ok(remove_floats::<Float64Type>(array, value)),
+ _ => {
+ let equal = spark_equality(array.values().as_ref(), value.as_ref())?;
+ remove_where(array, value, equal)
+ }
+ }
+}
+
+/// [`array_remove`] of the same value from every row, which a float array does in one pass over
+/// all the values, as DataFusion does for its own `array_remove`.
+fn remove_constant(array: &ListArray, needle: &ScalarValue) -> Result<ArrayRef> {
+ match (array.value_type(), needle) {
+ (DataType::Float32, ScalarValue::Float32(Some(needle))) => {
+ remove_float_constant::<Float32Type>(array, *needle)
+ }
+ (DataType::Float64, ScalarValue::Float64(Some(needle))) => {
+ remove_float_constant::<Float64Type>(array, *needle)
+ }
+ _ => array_remove(array, &needle.to_array_of_size(array.len())?),
+ }
+}
+
+fn remove_float_constant<T>(array: &ListArray, needle: T::Native) -> Result<ArrayRef>
+where
+ T: ArrowPrimitiveType,
+ T::Native: Float,
+{
+ let offsets = array.offsets();
+ let first = offsets[0] as usize;
+ let values = array
+ .values()
+ .slice(first, offsets[array.len()] as usize - first);
+ let floats = values.as_primitive::<T>();
+ let buffer = floats.values();
+ // For a fixed value, Spark's genEqual is one test per element: IEEE `!=`, under which the two
+ // zeros are equal and a NaN differs from every number, or `is_nan` for a NaN value.
+ let differs = if needle.is_nan() {
+ BooleanBuffer::collect_bool(buffer.len(), |index| !buffer[index].is_nan())
+ } else {
+ BooleanBuffer::collect_bool(buffer.len(), |index| buffer[index] != needle)
+ };
+ // Null elements stay.
+ let keep = match floats.nulls() {
+ Some(nulls) => &differs | &!nulls.inner(),
+ None => differs,
+ };
+ // A null row keeps its elements in the values; the list's null hides them.
+ let kept_offsets = kept_offsets(&keep, offsets);
+ let kept_values = filter(&values, &BooleanArray::new(keep, None))?;
+ Ok(with_values(
+ array,
+ kept_offsets,
+ kept_values,
+ array.nulls().cloned(),
+ ))
+}
+
+/// The offsets of the elements that `keep` keeps, for a `keep` that starts at `offsets[0]`: the
+/// number of its set bits before each offset, counted a word at a time in one pass.
+fn kept_offsets(keep: &BooleanBuffer, offsets: &OffsetBuffer<i32>) -> OffsetBuffer<i32> {
+ let chunks = keep.inner().bit_chunks(keep.offset(), keep.len());
+ let mut words = chunks.iter_padded();
+ let (mut word, mut word_start, mut before_word) = (words.next().unwrap_or(0), 0, 0);
+ let kept = offsets.iter().map(|offset| {
+ let end = (offset - offsets[0]) as usize;
+ while end >= word_start + 64 {
+ before_word += word.count_ones() as usize;
+ word = words.next().unwrap_or(0);
+ word_start += 64;
+ }
+ let below_end = word & ((1u64 << (end - word_start)) - 1);
+ (before_word + below_end.count_ones() as usize) as i32
+ });
+ OffsetBuffer::new(kept.collect::<Vec<i32>>().into())
+}
+
+/// Copies the floats of each row that are null or differ from the row's value in Spark's
+/// `genEqual`, straight from the flattened values. A row is null when the array or the value is.
+fn remove_floats<T>(array: &ListArray, value: &ArrayRef) -> ArrayRef
+where
+ T: ArrowPrimitiveType,
+ T::Native: Float,
+{
+ let values = array.values().as_primitive::<T>();
+ let needles = value.as_primitive::<T>().values();
+ let (buffer, nulls) = (values.values(), values.nulls());
+ let row_nulls = NullBuffer::union(array.nulls(), value.nulls());
+ let mut kept: Vec<T::Native> = Vec::with_capacity(values.len());
+ let mut validity = nulls.map(|_| BooleanBufferBuilder::new(values.len()));
+ let mut offsets = Vec::with_capacity(array.len() + 1);
+ offsets.push(0i32);
+ for (row, bounds) in array.offsets().windows(2).enumerate() {
+ if row_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(row)) {
+ let needle = needles[row];
+ let (start, end) = (bounds[0] as usize, bounds[1] as usize);
+ match (nulls, validity.as_mut()) {
+ (Some(nulls), Some(validity)) => {
+ for index in start..end {
+ let valid = nulls.is_valid(index);
+ if !valid || !compare_floats(buffer[index], needle).is_eq() {
+ kept.push(buffer[index]);
+ validity.append(valid);
+ }
+ }
+ }
+ _ => kept.extend(
+ buffer[start..end]
+ .iter()
+ .copied()
+ .filter(|&element| !compare_floats(element, needle).is_eq()),
+ ),
+ }
+ }
+ offsets.push(kept.len() as i32);
+ }
+ let kept = PrimitiveArray::<T>::new(
+ kept.into(),
+ validity.map(|mut validity| NullBuffer::new(validity.finish())),
+ );
+ with_values(
+ array,
+ OffsetBuffer::new(offsets.into()),
+ Arc::new(kept),
+ row_nulls,
+ )
+}
+
+/// Copies each row's elements except the non-null ones for which `equal(element, row)` holds. A
+/// row is null when the array or the value is.
+fn remove_where<F>(array: &ListArray, value: &ArrayRef, equal: F) -> Result<ArrayRef>
+where
+ F: Fn(usize, usize) -> bool,
+{
+ let values = array.values();
+ let element_nulls = values.logical_nulls();
+ let row_nulls = NullBuffer::union(array.nulls(), value.nulls());
+ let mut keep = BooleanBufferBuilder::new(values.len());
+ keep.append_n(values.len(), false);
+ let mut offsets = Vec::with_capacity(array.len() + 1);
+ offsets.push(0i32);
+ let mut kept = 0i32;
+ for (row, bounds) in array.offsets().windows(2).enumerate() {
+ if row_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(row)) {
+ for element in bounds[0] as usize..bounds[1] as usize {
+ let is_null = element_nulls
+ .as_ref()
+ .is_some_and(|nulls| nulls.is_null(element));
+ if is_null || !equal(element, row) {
+ keep.set_bit(element, true);
+ kept += 1;
+ }
+ }
+ }
+ offsets.push(kept);
+ }
+ let kept_values = filter(values, &BooleanArray::new(keep.finish(), None))?;
+ Ok(with_values(
+ array,
+ OffsetBuffer::new(offsets.into()),
+ kept_values,
+ row_nulls,
+ ))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::super::test_util::{bits, first_row_field_bits, list, EDGE_VALUES};
+ use super::*;
+ use crate::float_semantics::NEGATIVE_NAN;
+ use arrow::array::{Float32Array, Float64Array, StructArray};
+ use arrow::datatypes::{Field, Fields};
+ use datafusion::config::ConfigOptions;
+
+ fn invoke(array: ColumnarValue, value: ColumnarValue, rows: usize) -> Result<ColumnarValue> {
+ let return_field = Arc::new(Field::new("result", array.data_type(), true));
+ SparkArrayRemove::default().invoke_with_args(ScalarFunctionArgs {
+ args: vec![array, value],
+ arg_fields: vec![],
+ number_rows: rows,
+ return_field,
+ config_options: Arc::new(ConfigOptions::default()),
+ })
+ }
+
+ /// Spark's result for one row: drop the non-null elements that `genEqual` the value.
+ fn spark_remove(row: &[Option<f64>], value: f64) -> Vec<Option<u64>> {
+ row.iter()
+ .filter(|element| element.is_none_or(|e| !compare_floats(e, value).is_eq()))
+ .map(|element| element.map(f64::to_bits))
+ .collect()
+ }
+
+ /// Every edge value removed from an array of all edge values, so each zero and each NaN meets
+ /// every other, and the kept elements keep their bits.
+ #[test]
+ fn floats_follow_spark_equality() -> Result<()> {
+ let elements: Vec<Option<f64>> = EDGE_VALUES.to_vec();
+ let removable: Vec<f64> = EDGE_VALUES.iter().flatten().copied().collect();
+ let rows: Vec<Option<Vec<Option<f64>>>> =
+ removable.iter().map(|_| Some(elements.clone())).collect();
+ let result = invoke(
+ ColumnarValue::Array(list(&rows)),
+ ColumnarValue::Array(Arc::new(Float64Array::from(removable.clone()))),
+ rows.len(),
+ )?
+ .into_array(rows.len())?;
+ let expected: Vec<Option<Vec<Option<u64>>>> = removable
+ .iter()
+ .map(|&value| Some(spark_remove(&elements, value)))
+ .collect();
+ assert_eq!(bits(&result), expected);
+ Ok(())
+ }
+
+ /// Removing a constant takes one pass over all the values, which must still respect each row's
+ /// bounds, null rows and null elements, in a sliced array too.
+ #[test]
+ fn constant_value() -> Result<()> {
+ let rows = vec![
+ Some(vec![Some(1.0), Some(-0.0)]),
+ Some(vec![Some(0.0), None, Some(NEGATIVE_NAN)]),
+ None,
+ Some(vec![Some(f64::NAN), Some(0.0), Some(2.0)]),
+ ];
+ let full = list(&rows);
+ for (array, rows) in [
+ (Arc::clone(&full), &rows[..]),
+ (full.slice(1, 3), &rows[1..]),
+ ] {
+ for needle in [Some(0.0), Some(f64::NAN), None] {
+ let result = invoke(
+ ColumnarValue::Array(Arc::clone(&array)),
+ ColumnarValue::Scalar(ScalarValue::Float64(needle)),
+ rows.len(),
+ )?
+ .into_array(rows.len())?;
+ let expected: Vec<Option<Vec<Option<u64>>>> = rows
+ .iter()
+ .map(|row| {
+ let needle = needle?;
+ row.as_ref().map(|row| spark_remove(row, needle))
+ })
+ .collect();
+ assert_eq!(bits(&result), expected, "remove {needle:?}");
+ }
+ }
+ Ok(())
+ }
+
+ /// Kept counts are read from the mask a 64-bit word at a time, so these rows end on word
+ /// boundaries and inside words, span several words, or are empty or null, and so do those of
+ /// the slice, whose values start at a word boundary.
+ #[test]
+ fn constant_value_across_words() -> Result<()> {
+ let lengths = [64, 0, 64, 1, 63, 128, 5, 59, 130, 0, 70, 3];
+ let rows: Vec<Option<Vec<Option<f64>>>> = lengths
+ .iter()
+ .enumerate()
+ .map(|(row, &length)| {
+ (row != 7).then(|| {
+ (0..length)
+ .map(|i| EDGE_VALUES[(row + i) % EDGE_VALUES.len()])
+ .collect()
+ })
+ })
+ .collect();
+ let full = list(&rows);
+ for (array, rows) in [
+ (Arc::clone(&full), &rows[..]),
+ (full.slice(2, 8), &rows[2..10]),
+ ] {
+ for needle in [0.0, f64::NAN, 1.0] {
+ let result = invoke(
+ ColumnarValue::Array(Arc::clone(&array)),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(needle))),
+ rows.len(),
+ )?
+ .into_array(rows.len())?;
+ let expected: Vec<Option<Vec<Option<u64>>>> = rows
+ .iter()
+ .map(|row| row.as_ref().map(|row| spark_remove(row, needle)))
+ .collect();
+ assert_eq!(bits(&result), expected, "remove {needle}");
+ }
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn nulls() -> Result<()> {
+ let rows = vec![
+ Some(vec![Some(0.0), None, Some(-0.0)]),
+ None,
+ Some(vec![Some(1.0)]),
+ ];
+ let value = Float64Array::from(vec![Some(0.0), Some(0.0), None]);
+ let result = invoke(
+ ColumnarValue::Array(list(&rows)),
+ ColumnarValue::Array(Arc::new(value)),
+ 3,
+ )?
+ .into_array(3)?;
+ // The null element stays; a null array or a null value gives a null row.
+ assert_eq!(bits(&result), vec![Some(vec![None]), None, None]);
+ Ok(())
+ }
+
+ #[test]
+ fn float32_and_scalars() -> Result<()> {
+ let array = ScalarValue::List(Arc::new(
+ ListArray::from_iter_primitive::<Float32Type, _, _>(vec![Some(vec![
+ Some(-0.0f32),
+ Some(f32::from_bits(0xffc0_0000)),
+ Some(1.0),
+ ])]),
+ ));
+ let value = ScalarValue::Float32(Some(0.0));
+ let result = invoke(
+ ColumnarValue::Scalar(array.clone()),
+ ColumnarValue::Scalar(value),
+ 4,
+ )?;
+ let ColumnarValue::Scalar(ScalarValue::List(result)) = result else {
+ panic!("expected a list scalar");
+ };
+ let values = result.value(0);
+ let values: &Float32Array = values.as_primitive();
+ assert_eq!(values.len(), 2, "only the zero is removed");
+ let result = invoke(
+ ColumnarValue::Scalar(array),
+ ColumnarValue::Scalar(ScalarValue::Float32(Some(f32::NAN))),
+ 1,
+ )?;
+ let ColumnarValue::Scalar(ScalarValue::List(result)) = result else {
+ panic!("expected a list scalar");
+ };
+ let values = result.value(0);
+ let values: &Float32Array = values.as_primitive();
+ assert_eq!(values.values().as_ref(), &[-0.0, 1.0]);
+ assert_eq!(values.value(0).to_bits(), (-0.0f32).to_bits());
+ Ok(())
+ }
+
+ /// Structs compare their float fields in Spark's order too.
+ #[test]
+ fn nested_elements() -> Result<()> {
+ let fields = Fields::from(vec![Field::new("x", DataType::Float64, true)]);
+ let structs = |values: Vec<f64>| -> ArrayRef {
+ Arc::new(StructArray::new(
+ fields.clone(),
+ vec![Arc::new(Float64Array::from(values))],
+ None,
+ ))
+ };
+ let element = Field::new("item", DataType::Struct(fields.clone()), true);
+ let array = ListArray::new(
+ Arc::new(element),
+ OffsetBuffer::from_lengths([3]),
+ structs(vec![-0.0, NEGATIVE_NAN, 1.0]),
+ None,
+ );
+ for (value, expected) in [(0.0, vec![NEGATIVE_NAN, 1.0]), (f64::NAN, vec![-0.0, 1.0])] {
+ let result = invoke(
+ ColumnarValue::Array(Arc::new(array.clone())),
+ ColumnarValue::Array(structs(vec![value])),
+ 1,
+ )?
+ .into_array(1)?;
+ let expected: Vec<u64> = expected.iter().map(|v| v.to_bits()).collect();
+ assert_eq!(first_row_field_bits(&result), expected, "remove {value}");
+ }
+ Ok(())
+ }
+
+ /// Spark casts the value to the element type, so a value of another type is an error, not a
+ /// panic, whether it is a constant or a column.
+ #[test]
+ fn mismatched_value_type() {
+ let array = list(&[Some(vec![Some(1.0)])]);
+ for value in [
+ ColumnarValue::Scalar(ScalarValue::Float32(Some(1.0))),
+ ColumnarValue::Array(Arc::new(Float32Array::from(vec![1.0f32]))),
+ ] {
+ assert!(invoke(ColumnarValue::Array(Arc::clone(&array)), value, 1).is_err());
+ }
+ }
+}
diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs
index f5d87f09a8..ab40079036 100644
--- a/native/spark-expr/src/array_funcs/mod.rs
+++ b/native/spark-expr/src/array_funcs/mod.rs
@@ -18,6 +18,7 @@
mod array_extrema;
mod array_insert;
mod array_position;
+mod array_remove;
mod array_slice;
mod arrays_overlap;
mod arrays_zip;
@@ -28,10 +29,12 @@ mod list_positions;
mod nested_comparison;
mod sequence;
mod size;
+mod sort_array;
pub use array_extrema::SparkArrayExtrema;
pub use array_insert::ArrayInsert;
pub use array_position::SparkArrayPositionFunc;
+pub use array_remove::SparkArrayRemove;
pub use array_slice::SparkArraySlice;
pub use arrays_overlap::SparkArraysOverlap;
pub use arrays_zip::SparkArraysZipFunc;
@@ -42,3 +45,72 @@ pub use list_positions::ListPositionsExpr;
pub use nested_comparison::{spark_comparison, spark_in_list};
pub use sequence::spark_sequence;
pub use size::{spark_size, SparkSizeFunc};
+pub use sort_array::SparkSortArray;
+
+use arrow::array::{ArrayRef, ListArray};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use std::sync::Arc;
+
+/// A list with the field of `array` and the given row nulls, holding `values` at `offsets`.
+fn with_values(
+ array: &ListArray,
+ offsets: OffsetBuffer<i32>,
+ values: ArrayRef,
+ nulls: Option<NullBuffer>,
+) -> ArrayRef {
+ let (field, ..) = array.clone().into_parts();
+ Arc::new(ListArray::new(field, offsets, values, nulls))
+}
+
+#[cfg(test)]
+mod test_util {
+ use crate::float_semantics::{NEGATIVE_NAN, PAYLOAD_NAN};
+ use arrow::array::{ArrayRef, AsArray, ListArray};
+ use arrow::datatypes::Float64Type;
+ use std::sync::Arc;
+
+ /// The floats that Spark's rules treat specially, and a null.
+ pub(super) const EDGE_VALUES: [Option<f64>; 10] = [
+ Some(f64::NEG_INFINITY),
+ Some(-1.0),
+ Some(-0.0),
+ Some(0.0),
+ Some(1.0),
+ Some(f64::INFINITY),
+ Some(f64::NAN),
+ Some(NEGATIVE_NAN),
+ Some(PAYLOAD_NAN),
+ None,
+ ];
+
+ /// A list of `DOUBLE` rows.
+ pub(super) fn list(rows: &[Option<Vec<Option<f64>>>]) -> ArrayRef {
+ Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(
+ rows.iter().cloned(),
+ ))
+ }
+
+ /// The bits of every element of a list of `DOUBLE` rows, which tell `-0.0` from `0.0` and one
+ /// NaN from another.
+ pub(super) fn bits(array: &ArrayRef) -> Vec<Option<Vec<Option<u64>>>> {
+ array
+ .as_list::<i32>()
+ .iter()
+ .map(|row| {
+ row.map(|row| {
+ row.as_primitive::<Float64Type>()
+ .iter()
+ .map(|v| v.map(f64::to_bits))
+ .collect()
+ })
+ })
+ .collect()
+ }
+
+ /// The bits of the first field of the struct elements in the first row of a list.
+ pub(super) fn first_row_field_bits(array: &ArrayRef) -> Vec<u64> {
+ let row = array.as_list::<i32>().value(0);
+ let field = row.as_struct().column(0).as_primitive::<Float64Type>();
+ field.values().iter().map(|v| v.to_bits()).collect()
+ }
+}
diff --git a/native/spark-expr/src/array_funcs/sort_array.rs b/native/spark-expr/src/array_funcs/sort_array.rs
new file mode 100644
index 0000000000..a1b55fa5f7
--- /dev/null
+++ b/native/spark-expr/src/array_funcs/sort_array.rs
@@ -0,0 +1,521 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::cmp::Ordering;
+use std::sync::Arc;
+
+use super::with_values;
+use crate::float_semantics::{
+ compare_floats, compare_floats_java, float_gt, float_lt, spark_comparator,
+};
+use arrow::array::{
+ Array, ArrayRef, AsArray, BooleanBufferBuilder, ListArray, PrimitiveArray, UInt32Array,
+};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::compute::take;
+use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type};
+use datafusion::common::{exec_err, utils::take_function_args, Result, ScalarValue};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use num::Float;
+
+/// Spark's `sort_array` for arrays whose elements hold floats at any depth.
+///
+/// Spark sorts in its SQL ordering, in which `-0.0` equals `0.0` and all NaNs are equal and larger
+/// than every other value, at any depth, with a stable sort, so equal elements keep their order.
+/// Null elements come first when ascending and last when descending. DataFusion's `array_sort`
+/// uses IEEE 754 total order instead, in which a NaN with the sign bit set sorts first.
+///
+/// The arguments are the array, whether to sort ascending, and whether `-0.0` sorts before `0.0`.
+/// Spark's generated code sorts that way when it sorts an ascending array of `FLOAT` or `DOUBLE`
+/// that cannot hold a null with `java.util.Arrays.sort`. The serde decides, since it depends on
+/// Spark's `containsNull`, which Arrow's field nullability does not carry.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SparkSortArray {
+ signature: Signature,
+}
+
+impl Default for SparkSortArray {
+ fn default() -> Self {
+ Self {
+ signature: Signature::any(3, Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkSortArray {
+ fn name(&self) -> &str {
+ "spark_sort_array"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ Ok(arg_types[0].clone())
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ let [array, ascending, negative_zero_first] = take_function_args(self.name(), &args.args)?;
+ let order = match (constant(ascending), constant(negative_zero_first)) {
+ (Some(true), Some(false)) => Order::Ascending,
+ (Some(false), Some(false)) => Order::Descending,
+ (Some(true), Some(true)) => Order::Java,
+ _ => {
+ return exec_err!(
+ "spark_sort_array takes constant boolean flags, and puts -0.0 first only \
+ when sorting ascending"
+ )
+ }
+ };
+ let is_scalar = matches!(array, ColumnarValue::Scalar(_));
+ let rows = if is_scalar { 1 } else { args.number_rows };
+ let array = array.to_array(rows)?;
+ // Spark arrays use Arrow's 32-bit List layout.
+ let Some(list) = array.as_list_opt::<i32>() else {
+ return exec_err!("spark_sort_array takes a list, got {}", array.data_type());
+ };
+ let result = sort_array(list, order)?;
+ if is_scalar {
+ Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &result, 0,
+ )?))
+ } else {
+ Ok(ColumnarValue::Array(result))
+ }
+ }
+}
+
+fn constant(value: &ColumnarValue) -> Option<bool> {
+ match value {
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(value))) => Some(*value),
+ _ => None,
+ }
+}
+
+/// The order of a Spark `sort_array`.
+#[derive(Clone, Copy, PartialEq)]
+enum Order {
+ /// Spark's SQL ordering.
+ Ascending,
+ /// Spark's SQL ordering, reversed.
+ Descending,
+ /// [`compare_floats_java`], the order of `java.util.Arrays.sort`, which is ascending.
+ Java,
+}
+
+fn sort_array(array: &ListArray, order: Order) -> Result<ArrayRef> {
+ match array.value_type() {
+ DataType::Float32 => Ok(sort_floats::<Float32Type>(array, order)),
+ DataType::Float64 => Ok(sort_floats::<Float64Type>(array, order)),
+ _ => {
+ // The serde sorts -0.0 first only for FLOAT and DOUBLE elements.
+ let values = array.values();
+ let compare = spark_comparator(values.as_ref(), values.as_ref())?;
+ sort_by(array, order != Order::Descending, compare)
+ }
+ }
+}
+
+/// Sorts the floats of each row as Spark does, within one copy of all the values. Null elements go
+/// first when ascending and last when descending, as Spark's comparators put them.
+fn sort_floats<T>(array: &ListArray, order: Order) -> ArrayRef
+where
+ T: ArrowPrimitiveType,
+ T::Native: Float,
+{
+ let values = array.values().as_primitive::<T>();
+ let buffer = values.values();
+ let offsets = array.offsets();
+ let base = offsets[0] as usize;
+ let total = offsets[array.len()] as usize - base;
+ let rebased = offsets.clone().subtract(offsets[0]);
+ let Some(nulls) = values.nulls().filter(|nulls| nulls.null_count() > 0) else {
+ // Without null elements, each row sorts where it is.
+ let mut sorted = buffer[base..base + total].to_vec();
+ for (row, bounds) in offsets.windows(2).enumerate() {
+ if array.is_valid(row) {
+ let (start, end) = (bounds[0] as usize, bounds[1] as usize);
+ let original = buffer[start..end].iter().copied();
+ sort_row(&mut sorted[start - base..end - base], original, order);
+ }
+ }
+ let values = PrimitiveArray::<T>::new(sorted.into(), None);
+ return with_values(array, rebased, Arc::new(values), array.nulls().cloned());
+ };
+ // Each row's valid values are copied to its end when ascending, after its nulls, or to its
+ // start when descending, before them, and sorted there. The copy writes every element and
+ // moves on past only the valid ones. A null row's elements stay null, hidden by the list's
+ // null.
+ let mut sorted = vec![T::Native::default(); total];
+ let mut validity = BooleanBufferBuilder::new(total);
+ for (row, bounds) in offsets.windows(2).enumerate() {
+ let (start, end) = (bounds[0] as usize, bounds[1] as usize);
+ if array.is_null(row) {
+ validity.append_n(end - start, false);
+ continue;
+ }
+ let (row_start, row_end) = (start - base, end - base);
+ let valid = if order == Order::Descending {
+ let mut next = row_start;
+ for index in start..end {
+ sorted[next] = buffer[index];
+ next += usize::from(nulls.is_valid(index));
+ }
+ row_start..next
+ } else {
+ let mut next = row_end;
+ for index in (start..end).rev() {
+ sorted[next - 1] = buffer[index];
+ next -= usize::from(nulls.is_valid(index));
+ }
+ next..row_end
+ };
+ validity.append_n(valid.start - row_start, false);
+ validity.append_n(valid.len(), true);
+ validity.append_n(row_end - valid.end, false);
+ let original = (start..end)
+ .filter(|&index| nulls.is_valid(index))
+ .map(|index| buffer[index]);
+ sort_row(&mut sorted[valid], original, order);
+ }
+ let values = PrimitiveArray::<T>::new(sorted.into(), Some(NullBuffer::new(validity.finish())));
+ with_values(array, rebased, Arc::new(values), array.nulls().cloned())
+}
+
+/// The longest row that [`sort_row`] sorts stably. The standard library sorts rows this short by
+/// insertion, stable or not, so an unstable sort would only add [`restore_ties`].
+const STABLE_SORT_MAX_LEN: usize = 20;
+
+/// Sorts a row of floats to the result of Spark's stable sort, given `original`, the row's values
+/// in their original order. A longer row sorts faster with an unstable sort, after which
+/// [`restore_ties`] puts the tied elements back in their original order.
+fn sort_row<N: Float>(row: &mut [N], original: impl Iterator<Item = N> + Clone, order: Order) {
+ let stable = row.len() <= STABLE_SORT_MAX_LEN;
+ sort_by_order(row, stable, order);
+ if !stable {
+ restore_ties(row, original, order);
+ }
+}
+
+/// Sorts a row in `order`, with each comparator compiled separately.
+fn sort_by_order<N: Float>(row: &mut [N], stable: bool, order: Order) {
+ match (order, stable) {
+ (Order::Ascending, true) => row.sort_by(|a, b| compare_floats(*a, *b)),
+ (Order::Ascending, false) => row.sort_unstable_by(|a, b| compare_floats(*a, *b)),
+ (Order::Descending, true) => row.sort_by(|a, b| compare_floats(*b, *a)),
+ (Order::Descending, false) => row.sort_unstable_by(|a, b| compare_floats(*b, *a)),
+ (Order::Java, true) => row.sort_by(|a, b| compare_floats_java(*a, *b)),
+ (Order::Java, false) => row.sort_unstable_by(|a, b| compare_floats_java(*a, *b)),
+ }
+}
+
+/// Puts the tied elements of a row sorted by an unstable sort back in their order in `original`,
+/// the row's values in their original order, as Spark's stable sort leaves them.
+///
+/// The only tied elements whose bits can differ are NaNs, which sort last unless descending, and
+/// zeros of either sign, except in [`Order::Java`]. Each forms one run.
+fn restore_ties<N: Float>(row: &mut [N], original: impl Iterator<Item = N> + Clone, order: Order) {
+ let len = row.len();
+ let nans = if order == Order::Descending {
+ 0..row.iter().take_while(|v| v.is_nan()).count()
+ } else {
+ len - row.iter().rev().take_while(|v| v.is_nan()).count()..len
+ };
+ restore_order(&mut row[nans], original.clone().filter(|v| v.is_nan()));
+ let zero = N::zero();
+ let first = match order {
+ Order::Ascending => row.partition_point(|v| float_lt(*v, zero)),
+ Order::Descending => row.partition_point(|v| float_gt(*v, zero)),
+ Order::Java => return,
+ };
+ let end = first + row[first..].iter().take_while(|v| **v == zero).count();
+ restore_order(&mut row[first..end], original.filter(|v| *v == zero));
+}
+
+/// Overwrites a run of tied elements with `values`, the same elements in their original order.
+fn restore_order<N>(run: &mut [N], values: impl Iterator<Item = N>) {
+ if run.len() > 1 {
+ run.iter_mut()
+ .zip(values)
+ .for_each(|(slot, value)| *slot = value);
+ }
+}
+
+/// Sorts the elements of each row stably with `compare`, reversed when descending. `compare` puts
+/// null elements first, as Spark's ascending comparator does, so reversing it puts them last, as
+/// the descending one does.
+fn sort_by<F>(array: &ListArray, ascending: bool, compare: F) -> Result<ArrayRef>
+where
+ F: Fn(usize, usize) -> Ordering,
+{
+ let values = array.values();
+ if values.len() > u32::MAX as usize {
+ return exec_err!("sort_array cannot sort more than {} elements", u32::MAX);
+ }
+ let mut indices: Vec<u32> = Vec::with_capacity(values.len());
+ let mut offsets = Vec::with_capacity(array.len() + 1);
+ offsets.push(0i32);
+ for (row, bounds) in array.offsets().windows(2).enumerate() {
+ if array.is_valid(row) {
+ let start = indices.len();
+ indices.extend(bounds[0] as u32..bounds[1] as u32);
+ indices[start..].sort_by(|&a, &b| {
+ let (a, b) = (a as usize, b as usize);
+ if ascending {
+ compare(a, b)
+ } else {
+ compare(b, a)
+ }
+ });
+ }
+ offsets.push(indices.len() as i32);
+ }
+ let sorted = take(values.as_ref(), &UInt32Array::from(indices), None)?;
+ Ok(with_values(
+ array,
+ OffsetBuffer::new(offsets.into()),
+ sorted,
+ array.nulls().cloned(),
+ ))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::super::test_util::{bits, first_row_field_bits, list, EDGE_VALUES};
+ use super::*;
+ use crate::float_semantics::NEGATIVE_NAN;
+ use arrow::array::{Float32Array, Float64Array, StructArray};
+ use arrow::datatypes::{Field, Fields};
+ use datafusion::config::ConfigOptions;
+
+ fn invoke(array: ArrayRef, ascending: bool, negative_zero_first: bool) -> Result<ArrayRef> {
+ let rows = array.len();
+ let return_field = Arc::new(Field::new("result", array.data_type().clone(), true));
+ SparkSortArray::default()
+ .invoke_with_args(ScalarFunctionArgs {
+ args: vec![
+ ColumnarValue::Array(array),
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(ascending))),
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(negative_zero_first))),
+ ],
+ arg_fields: vec![],
+ number_rows: rows,
+ return_field,
+ config_options: Arc::new(ConfigOptions::default()),
+ })?
+ .into_array(rows)
+ }
+
+ fn row_bits(row: &[Option<f64>]) -> Vec<Option<u64>> {
+ row.iter().map(|v| v.map(f64::to_bits)).collect()
+ }
+
+ /// Pseudo-random sequences of edge values, the same on every run, of up to 82 elements, since
+ /// sorts longer than 20 elements are where an unstable sort reorders ties.
+ fn sequences() -> Vec<Option<Vec<Option<f64>>>> {
+ let mut state = 0x2545_f491_4f6c_dd1du64;
+ let mut next = move || {
+ state ^= state << 13;
+ state ^= state >> 7;
+ state ^= state << 17;
+ state
+ };
+ (0..200)
+ .map(|i| {
+ Some(
+ (0..i * 7 % 83)
+ .map(|_| EDGE_VALUES[(next() % EDGE_VALUES.len() as u64) as usize])
+ .collect(),
+ )
+ })
+ .chain([None])
+ .collect()
+ }
+
+ /// The sequences with their null elements removed.
+ fn without_nulls(rows: &[Option<Vec<Option<f64>>>]) -> Vec<Option<Vec<Option<f64>>>> {
+ rows.iter()
+ .map(|row| {
+ row.as_ref()
+ .map(|row| row.iter().flatten().map(|v| Some(*v)).collect())
+ })
+ .collect()
+ }
+
+ /// Spark's comparator sort: nulls first when ascending and last when descending, then
+ /// `compareDoubles`, reversed when descending, in a stable sort.
+ fn spark_sort(row: &[Option<f64>], ascending: bool) -> Vec<Option<u64>> {
+ let mut sorted = row.to_vec();
+ sorted.sort_by(|a, b| match (a, b) {
+ (None, None) => Ordering::Equal,
+ (None, Some(_)) if ascending => Ordering::Less,
+ (None, Some(_)) => Ordering::Greater,
+ (Some(_), None) if ascending => Ordering::Greater,
+ (Some(_), None) => Ordering::Less,
+ (Some(a), Some(b)) if ascending => compare_floats(*a, *b),
+ (Some(a), Some(b)) => compare_floats(*b, *a),
+ });
+ row_bits(&sorted)
+ }
+
+ /// An array that can hold nulls sorts stably in Spark's order: the zeros tie and keep their
+ /// order, and so do the NaNs, which sort last. That holds with and without null elements, and
+ /// in a sliced array, whose values start past its first offset.
+ #[test]
+ fn floats_follow_spark_order() -> Result<()> {
+ let with_nulls = sequences();
+ let without_nulls = without_nulls(&with_nulls);
+ for rows in [with_nulls, without_nulls] {
+ for ascending in [true, false] {
+ let expected: Vec<Option<Vec<Option<u64>>>> = rows
+ .iter()
+ .map(|row| row.as_ref().map(|row| spark_sort(row, ascending)))
+ .collect();
+ let result = invoke(list(&rows), ascending, false)?;
+ assert_eq!(bits(&result), expected, "ascending={ascending}");
+ let result = invoke(list(&rows).slice(37, 120), ascending, false)?;
+ assert_eq!(bits(&result), &expected[37..157], "ascending={ascending}");
+ }
+ }
+ Ok(())
+ }
+
+ /// `java.util.Arrays.sort` on a `double[]` orders by `Double.compare`, which compares
+ /// numerically and then by `doubleToLongBits`, so that all NaNs are equal and -0.0 comes before
+ /// 0.0, and it keeps the NaNs in their original order.
+ #[test]
+ fn negative_zero_first_follows_java() -> Result<()> {
+ let long_bits = |v: f64| {
+ if v.is_nan() {
+ f64::NAN.to_bits() as i64
+ } else {
+ v.to_bits() as i64
+ }
+ };
+ let rows = without_nulls(&sequences());
+ let expected: Vec<Option<Vec<Option<u64>>>> = rows
+ .iter()
+ .map(|row| {
+ row.as_ref().map(|row| {
+ let mut sorted: Vec<f64> = row.iter().flatten().copied().collect();
+ sorted.sort_by(|a, b| {
+ a.partial_cmp(b)
+ .filter(|ordering| ordering.is_ne())
+ .unwrap_or_else(|| long_bits(*a).cmp(&long_bits(*b)))
+ });
+ sorted.iter().map(|v| Some(v.to_bits())).collect()
+ })
+ })
+ .collect();
+ let result = invoke(list(&rows), true, true)?;
+ assert_eq!(bits(&result), expected);
+ Ok(())
+ }
+
+ /// With `-0.0` first, the zeros no longer tie and keep their order.
+ #[test]
+ fn negative_zero_first() -> Result<()> {
+ let rows = vec![Some(vec![
+ Some(0.0),
+ Some(1.0),
+ Some(-0.0),
+ Some(NEGATIVE_NAN),
+ ])];
+ let cases = [
+ (true, vec![-0.0, 0.0, 1.0, NEGATIVE_NAN]),
+ (false, vec![0.0, -0.0, 1.0, NEGATIVE_NAN]),
+ ];
+ for (negative_zero_first, expected) in cases {
+ let result = invoke(list(&rows), true, negative_zero_first)?;
+ let expected: Vec<Option<u64>> = expected.iter().map(|v| Some(v.to_bits())).collect();
+ assert_eq!(
+ bits(&result),
+ vec![Some(expected)],
+ "negative_zero_first={negative_zero_first}"
+ );
+ }
+ Ok(())
+ }
+
+ /// A non-list, and `-0.0` first in a descending sort, which Spark never asks for.
+ #[test]
+ fn invalid_arguments() {
+ assert!(invoke(Arc::new(Float64Array::from(vec![1.0])), true, false).is_err());
+ assert!(invoke(list(&[Some(vec![Some(0.0)])]), false, true).is_err());
+ }
+
+ #[test]
+ fn float32() -> Result<()> {
+ let array = Arc::new(ListArray::from_iter_primitive::<Float32Type, _, _>(vec![
+ Some(vec![
+ Some(f32::NAN),
+ Some(0.0f32),
+ Some(f32::from_bits(0xffc0_0000)),
+ Some(-0.0),
+ None,
+ ]),
+ ]));
+ let result = invoke(array, true, false)?;
+ let row = result.as_list::<i32>().value(0);
+ let row: &Float32Array = row.as_primitive();
+ let actual: Vec<Option<u32>> = row.iter().map(|v| v.map(f32::to_bits)).collect();
+ let expected = [
+ None,
+ Some(0.0f32),
+ Some(-0.0),
+ Some(f32::NAN),
+ Some(f32::from_bits(0xffc0_0000)),
+ ];
+ let expected: Vec<Option<u32>> = expected.iter().map(|v| v.map(f32::to_bits)).collect();
+ assert_eq!(actual, expected);
+ Ok(())
+ }
+
+ /// Struct elements compare their float fields in Spark's order and also sort stably.
+ #[test]
+ fn nested_elements() -> Result<()> {
+ let fields = Fields::from(vec![Field::new("x", DataType::Float64, true)]);
+ let element = Field::new("item", DataType::Struct(fields.clone()), true);
+ let values = vec![1.0, NEGATIVE_NAN, 0.0, f64::NAN, -0.0, -1.0];
+ let array: ArrayRef = Arc::new(ListArray::new(
+ Arc::new(element),
+ OffsetBuffer::from_lengths([values.len()]),
+ Arc::new(StructArray::new(
+ fields,
+ vec![Arc::new(Float64Array::from(values))],
+ None,
+ )),
+ None,
+ ));
+ for (ascending, expected) in [
+ (true, vec![-1.0, 0.0, -0.0, 1.0, NEGATIVE_NAN, f64::NAN]),
+ (false, vec![NEGATIVE_NAN, f64::NAN, 1.0, 0.0, -0.0, -1.0]),
+ ] {
+ let result = invoke(Arc::clone(&array), ascending, false)?;
+ let expected: Vec<u64> = expected.iter().map(|v| v.to_bits()).collect();
+ assert_eq!(
+ first_row_field_bits(&result),
+ expected,
+ "ascending={ascending}"
+ );
+ }
+ Ok(())
+ }
+}
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs
index 6ccaff2b2d..7b3c980639 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -30,11 +30,11 @@ use crate::{
spark_ceil, spark_day_name, spark_decimal_div, spark_decimal_integral_div, spark_floor,
spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding,
spark_round, spark_rpad, spark_sequence, spark_to_time, spark_unhex, spark_unscaled_value,
- EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap,
- SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek,
- SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate,
- SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMapExtract, SparkNextDay,
- SparkSecondsToTimestamp, SparkSizeFunc, SparkWeekDay,
+ EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, SparkArrayRemove, SparkArraySlice,
+ SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc,
+ SparkDayOfWeek, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform,
+ SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMapExtract,
+ SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, SparkSortArray, SparkWeekDay,
};
use arrow::datatypes::DataType;
use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -330,8 +330,10 @@ fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> {
Arc::new(ScalarUDF::new_from_impl(SparkArrayExtrema::new(true))),
Arc::new(ScalarUDF::new_from_impl(SparkArrayExtrema::new(false))),
Arc::new(ScalarUDF::new_from_impl(SparkArrayPositionFunc::default())),
+ Arc::new(ScalarUDF::new_from_impl(SparkArrayRemove::default())),
Arc::new(ScalarUDF::new_from_impl(SparkArraySlice::default())),
Arc::new(ScalarUDF::new_from_impl(SparkArraysOverlap::default())),
+ Arc::new(ScalarUDF::new_from_impl(SparkSortArray::default())),
Arc::new(ScalarUDF::new_from_impl(SparkContains::default())),
Arc::new(ScalarUDF::new_from_impl(SparkDateDiff::default())),
Arc::new(ScalarUDF::new_from_impl(SparkDateFromUnixDate::default())),
diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs
index 7d79c3dd31..b4ab1ae130 100644
--- a/native/spark-expr/src/float_semantics/mod.rs
+++ b/native/spark-expr/src/float_semantics/mod.rs
@@ -29,6 +29,10 @@
//! [`normalize_floats`] and [`normalize_nested_floats`].
//! - `java.lang.Double.equals`, which boxed keys such as those in `OpenHashSet` use: all NaNs are
//! equal, but `-0.0` and `0.0` are distinct. See [`canonicalize_nan`].
+//! - `java.lang.Double.compare`, which `java.util.Arrays.sort` on a primitive array uses: the SQL
+//! ordering, except that `-0.0` sorts before `0.0`. Spark's generated code sorts an ascending
+//! `sort_array` of `FLOAT` or `DOUBLE` elements that cannot be null this way. See
+//! [`compare_floats_java`].
//! - `Murmur3Hash` and `XxHash64`: `-0.0` hashes as `0.0`, and NaN as the canonical NaN. See
//! [`hash_input`].
//!
@@ -91,6 +95,16 @@ pub fn compare_floats<T: Float>(left: T, right: T) -> Ordering {
}
}
+/// `java.lang.Double.compare`: Spark's SQL ordering, except that `-0.0` sorts before `0.0`.
+#[inline]
+pub fn compare_floats_java<T: Float>(left: T, right: T) -> Ordering {
+ match compare_floats(left, right) {
+ // Equal values that are not NaN differ in sign only when they are -0.0 and 0.0.
+ Ordering::Equal if !left.is_nan() => right.is_sign_negative().cmp(&left.is_sign_negative()),
+ ordering => ordering,
+ }
+}
+
/// Whether `left` sorts before `right` in Spark's SQL ordering, the same as
/// `compare_floats(left, right).is_lt()`. As a single test it compiles to a well-predicted branch
/// in a scan for a minimum, where the three-way comparison is several times slower.
@@ -123,10 +137,10 @@ pub fn hash_input<T: Float>(v: T) -> T {
/// A NaN with the sign bit set, which arithmetic produces on x86-64.
#[cfg(test)]
-const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000);
+pub(crate) const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000);
/// A signaling NaN with a payload.
#[cfg(test)]
-const PAYLOAD_NAN: f64 = f64::from_bits(0x7ff0_0000_0000_0001);
+pub(crate) const PAYLOAD_NAN: f64 = f64::from_bits(0x7ff0_0000_0000_0001);
#[cfg(test)]
mod tests {
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 6cfe5bc92b..11ae8ac041 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -1088,10 +1088,10 @@ object CometConf extends ShimCometConf {
.category(CATEGORY_EXEC)
.doc(
"When enabled, fall back to Spark for floating-point operations that may differ from " +
- "Spark, such as comparing -0.0 and 0.0, sorting floating-point values nested in " +
- "arrays, structs, or maps, or sorting the elements of a floating-point array with " +
- "`sort_array`. Scalar `ORDER BY`, window ordering and range partitioning keys are " +
- "unaffected, because Comet normalizes those comparison keys to match Spark. " +
+ "Spark, such as comparing -0.0 and 0.0, or sorting floating-point values nested in " +
+ "arrays, structs, or maps. Scalar `ORDER BY`, window ordering and range partitioning " +
+ "keys are unaffected, because Comet normalizes those comparison keys to match Spark, " +
+ "and so is `sort_array`, which follows Spark's ordering. " +
s"$COMPAT_GUIDE.")
.booleanConf
.createWithDefault(false)
diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
index 8cdabebb85..6a5ce01561 100644
--- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
@@ -27,7 +27,6 @@ import org.apache.spark.sql.catalyst.util.GenericArrayData
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
-import org.apache.comet.CometConf
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
import org.apache.comet.DataTypeSupport.{deepNullable, isComplexType}
import org.apache.comet.serde.QueryPlanSerde._
@@ -46,8 +45,16 @@ object CometArrayRemove
binding: Boolean): Option[ExprOuterClass.Expr] = {
val arrayExprProto = exprToProtoInternal(expr.left, inputs, binding)
val keyExprProto = exprToProtoInternal(expr.right, inputs, binding)
-
- scalarFunctionExprToProto("array_remove_all", arrayExprProto, keyExprProto)
+ // DataFusion's array_remove_all compares floats by their bits. spark_array_remove compares
+ // them as Spark does, with -0.0 equal to 0.0 and all NaNs equal.
+ val elementType = expr.left.dataType.asInstanceOf[ArrayType].elementType
+ val function =
+ if (SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])) {
+ "spark_array_remove"
+ } else {
+ "array_remove_all"
+ }
+ scalarFunctionExprToProto(function, arrayExprProto, keyExprProto)
}
}
@@ -140,10 +147,6 @@ object CometArrayContains
object CometSortArray extends CometExpressionSerde[SortArray] with CodegenDispatchFallback {
- override def getIncompatibleReasons(): Seq[String] = Seq(
- "When `" + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key + "=true`, sorting on" +
- " floating-point types is not 100% compatible with Spark")
-
private def supportedSortArrayElementType(dt: DataType): Boolean = {
dt match {
case _: NullType =>
@@ -163,16 +166,13 @@ object CometSortArray extends CometExpressionSerde[SortArray] with CodegenDispat
if (!supportedSortArrayElementType(elementType)) {
Unsupported(Some(s"Sort on array element type $elementType is not supported"))
} else {
- SupportLevel
- .strictFloatingPointReason(elementType, "Sorting on floating-point")
- .map(reason => Incompatible(Some(reason)))
- .getOrElse(expr.ascendingOrder match {
- // Spark 3.x requires a boolean Literal; Spark 4.0+ widens ascendingOrder to any
- // foldable boolean. Accept both; convert evaluates the foldable expression.
- case ao if ao.foldable && ao.dataType == BooleanType => Compatible()
- case other =>
- Unsupported(Some(s"ascendingOrder must be a foldable boolean: $other"))
- })
+ expr.ascendingOrder match {
+ // Spark 3.x requires a boolean Literal; Spark 4.0+ widens ascendingOrder to any
+ // foldable boolean. Accept both; convert evaluates the foldable expression.
+ case ao if ao.foldable && ao.dataType == BooleanType => Compatible()
+ case other =>
+ Unsupported(Some(s"ascendingOrder must be a foldable boolean: $other"))
+ }
}
}
@@ -184,18 +184,30 @@ object CometSortArray extends CometExpressionSerde[SortArray] with CodegenDispat
// ascendingOrder is a foldable boolean (gated in getSupportLevel). Evaluate it; a null result
// unboxes to false, matching Spark's `right.eval().asInstanceOf[Boolean]`.
val ascending = expr.ascendingOrder.eval(EmptyRow).asInstanceOf[Boolean]
- val direction = if (ascending) "ASC" else "DESC"
- val nullOrdering = if (ascending) "NULLS FIRST" else "NULLS LAST"
- val sortDirectionExprProto = exprToProtoInternal(Literal(direction), inputs, binding)
- val nullOrderingExprProto = exprToProtoInternal(Literal(nullOrdering), inputs, binding)
-
- val sortArrayScalarExpr =
+ val arrayType = expr.base.dataType.asInstanceOf[ArrayType]
+ if (SupportLevel.containsType(
+ arrayType.elementType,
+ classOf[FloatType],
+ classOf[DoubleType])) {
+ // DataFusion's array_sort orders floats by IEEE 754 total order. spark_sort_array follows
+ // Spark's ordering, except that -0.0 sorts before 0.0 where Spark's generated code sorts
+ // with java.util.Arrays.sort: ascending, over FLOAT or DOUBLE elements that cannot be null.
+ val negativeZeroFirst = ascending && !arrayType.containsNull &&
+ (arrayType.elementType == FloatType || arrayType.elementType == DoubleType)
+ scalarFunctionExprToProto(
+ "spark_sort_array",
+ arrayExprProto,
+ exprToProtoInternal(Literal(ascending), inputs, binding),
+ exprToProtoInternal(Literal(negativeZeroFirst), inputs, binding))
+ } else {
+ val direction = if (ascending) "ASC" else "DESC"
+ val nullOrdering = if (ascending) "NULLS FIRST" else "NULLS LAST"
scalarFunctionExprToProto(
"array_sort",
arrayExprProto,
- sortDirectionExprProto,
- nullOrderingExprProto)
- sortArrayScalarExpr
+ exprToProtoInternal(Literal(direction), inputs, binding),
+ exprToProtoInternal(Literal(nullOrdering), inputs, binding))
+ }
}
}
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_remove_floating_point.sql b/spark/src/test/resources/sql-tests/expressions/array/array_remove_floating_point.sql
new file mode 100644
index 0000000000..d944d837f4
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_remove_floating_point.sql
@@ -0,0 +1,63 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- array_remove compares floats the way Spark's generated code does: -0.0 equals 0.0, and every
+-- NaN equals every other NaN, at any depth of an array or struct element. Null elements stay, and
+-- a null array or value gives null.
+--
+-- `-d` flips the sign bit, so `-d` of the NaN row is a NaN with the sign bit set on every
+-- platform. That is the NaN that arithmetic produces on x86-64.
+
+statement
+CREATE TABLE ar_float(id INT, d DOUBLE, f FLOAT) USING parquet
+
+statement
+INSERT INTO ar_float VALUES
+ (1, 0.0D, float('0.0')),
+ (2, double('-0.0'), float('-0.0')),
+ (3, double('NaN'), float('NaN')),
+ (4, 1.0D, float('1.0')),
+ (5, NULL, NULL)
+
+-- Removing a zero removes both zeros, and removing a NaN removes both NaNs
+query
+SELECT id, array_remove(array(d, -d, 1.0D, NULL), 0.0D), array_remove(array(d, -d, 1.0D), d),
+ array_remove(array(d, -d, 1.0D), -d)
+FROM ar_float
+
+query
+SELECT id, array_remove(array(f, -f, 1.0F, NULL), 0.0F), array_remove(array(f, -f), -f)
+FROM ar_float
+
+-- A literal array with a value from a column, and a null array
+query
+SELECT id, array_remove(array(0.0D, -0.0D, double('NaN'), 1.0D), d),
+ array_remove(IF(d IS NULL, NULL, array(d, -d, 1.0D)), 0.0D)
+FROM ar_float
+
+-- Arrays as elements compare their floats the same way. Struct elements fall back to Spark
+-- (#1307).
+query
+SELECT id, array_remove(array(array(d), array(-d), array(1.0D)), array(d)),
+ array_remove(array(array(-d, 1.0D), array(d, 1.0D)), array(-d, 1.0D))
+FROM ar_float
+
+-- Only literals, which the suite does not constant fold. `-0.0D` is a literal with the sign bit
+-- set, while `double('NaN')` is a cast.
+query
+SELECT array_remove(array(0.0D, -0.0D, 1.0D), -0.0D), array_remove(array(-0.0D, 0.0D), 0.0D),
+ array_remove(array(double('NaN'), 1.0D), double('NaN'))
diff --git a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql
index ece58525d4..a6e9499c04 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql
@@ -47,7 +47,8 @@ SELECT reverse(a) FROM routing_arrays
query expect_fallback(array_contains: spark.comet.exec.scalaUDF.codegen.enabled=false)
SELECT array_contains(f, 1.0D) FROM routing_arrays
-query expect_fallback(sort_array: spark.comet.exec.scalaUDF.codegen.enabled=false)
+-- sort_array on floats follows Spark natively, even in strict floating-point mode.
+query expect_native(sort_array)
SELECT sort_array(f) FROM routing_arrays
query expect_fallback(sort_array: spark.comet.exec.scalaUDF.codegen.enabled=false)
diff --git a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql
index 03d4d3120a..1e688857a9 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql
@@ -47,7 +47,8 @@ SELECT reverse(a) FROM routing_arrays
query expect_dispatch(array_contains)
SELECT array_contains(f, 1.0D) FROM routing_arrays
-query expect_dispatch(sort_array)
+-- sort_array on floats follows Spark natively, even in strict floating-point mode.
+query expect_native(sort_array)
SELECT sort_array(f) FROM routing_arrays
query expect_dispatch(sort_array)
diff --git a/spark/src/test/resources/sql-tests/expressions/array/sort_array_floating_point.sql b/spark/src/test/resources/sql-tests/expressions/array/sort_array_floating_point.sql
new file mode 100644
index 0000000000..95a6f6a8d7
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/array/sort_array_floating_point.sql
@@ -0,0 +1,68 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- sort_array orders floats the way Spark does, which depends on whether the elements can be null.
+-- When they can, or when sorting descending, Spark compares in its SQL ordering with a stable
+-- sort: -0.0 and 0.0 tie and keep their order, and every NaN ties with every other NaN and sorts
+-- after every other value. When an ascending array's elements cannot be null, Spark's generated
+-- code calls java.util.Arrays.sort instead, which puts -0.0 before 0.0.
+--
+-- `-d` flips the sign bit, so `-d` of the NaN row is a NaN with the sign bit set on every
+-- platform. That is the NaN that arithmetic produces on x86-64.
+
+statement
+CREATE TABLE sa_float(id INT, d DOUBLE, f FLOAT) USING parquet
+
+statement
+INSERT INTO sa_float VALUES
+ (1, 0.0D, float('0.0')),
+ (2, double('-0.0'), float('-0.0')),
+ (3, double('NaN'), float('NaN')),
+ (4, 1.0D, float('1.0')),
+ (5, NULL, NULL),
+ (6, double('-Infinity'), float('-Infinity'))
+
+-- Elements that can be null, with both zeros and both NaNs in both orders
+query
+SELECT id, sort_array(array(d, -d, 1.0D)), sort_array(array(-d, d, 1.0D)),
+ sort_array(array(d, -d, 1.0D), false), sort_array(array(-d, d, NULL), false)
+FROM sa_float
+
+query
+SELECT id, sort_array(array(f, -f, 1.0F, NULL)), sort_array(array(-f, f), false)
+FROM sa_float
+
+-- A null array
+query
+SELECT id, sort_array(IF(d IS NULL, NULL, array(d, -d, 1.0D))) FROM sa_float
+
+-- Elements that cannot be null
+query
+SELECT id, sort_array(array(coalesce(d, 0.0D), coalesce(-d, 0.0D), 1.0D)),
+ sort_array(array(coalesce(-d, 0.0D), coalesce(d, 0.0D)), false),
+ sort_array(array(coalesce(f, 0.0F), coalesce(-f, 0.0F)))
+FROM sa_float
+
+query
+SELECT sort_array(array(0.0D, -0.0D, 1.0D)), sort_array(array(0.0F, -0.0F)),
+ sort_array(array(0.0D, -0.0D), false)
+
+-- Arrays and structs as elements compare their floats in Spark's order, with a stable sort
+query
+SELECT id, sort_array(array(array(d), array(-d), array(1.0D))),
+ sort_array(array(named_struct('x', d), named_struct('x', -d)), false)
+FROM sa_float
diff --git a/spark/src/test/resources/sql-tests/expressions/array/sort_array_strict_fp.sql b/spark/src/test/resources/sql-tests/expressions/array/sort_array_strict_fp.sql
index db3453bd23..2b0a3aab2b 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/sort_array_strict_fp.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/sort_array_strict_fp.sql
@@ -15,9 +15,8 @@
-- specific language governing permissions and limitations
-- under the License.
--- With strict floating-point mode on, Comet's native float sort is Incompatible with Spark
--- (NaN / signed-zero ordering). CometSortArray mixes in CodegenDispatchFallback, so it routes
--- through the JVM codegen dispatcher and matches Spark exactly instead of falling back.
+-- Comet's native sort_array follows Spark's float ordering, so strict floating-point mode keeps it
+-- native. sort_array_floating_point.sql covers the ordering itself.
-- Config: spark.comet.exec.strictFloatingPoint=true
@@ -52,6 +51,6 @@ SELECT sort_array(arr) FROM test_sort_array_strict_float
query
SELECT sort_array(arr, false) FROM test_sort_array_strict_float
--- nested float arrays also dispatch
+-- an array built from literals
query
SELECT sort_array(array(CAST('NaN' AS DOUBLE), double('-0.0'), 0.0, 1.0, CAST('-Infinity' AS DOUBLE)))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org