From 32970e03df665529bc8924cb7aaaade7abf5cb07 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 1 Oct 2026 11:15:37 -0600 Subject: [PATCH 1/4] 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. --- .../expression-audits/array_funcs.md | 8 +- .../latest/compatibility/floating-point.md | 12 + docs/source/user-guide/latest/expressions.md | 2 +- .../user-guide/latest/tuning/operators.md | 5 +- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/float_arrays.rs | 147 +++++ .../src/array_funcs/array_remove.rs | 522 ++++++++++++++++ native/spark-expr/src/array_funcs/mod.rs | 4 + .../spark-expr/src/array_funcs/sort_array.rs | 560 ++++++++++++++++++ native/spark-expr/src/comet_scalar_funcs.rs | 12 +- .../scala/org/apache/comet/CometConf.scala | 8 +- .../scala/org/apache/comet/serde/arrays.scala | 50 +- .../array/array_remove_floating_point.sql | 59 ++ .../array/routing_arrays_disabled.sql | 3 +- .../array/routing_arrays_enabled.sql | 3 +- .../array/sort_array_floating_point.sql | 68 +++ .../array/sort_array_strict_fp.sql | 7 +- 17 files changed, 1428 insertions(+), 46 deletions(-) create mode 100644 native/spark-expr/benches/float_arrays.rs create mode 100644 native/spark-expr/src/array_funcs/array_remove.rs create mode 100644 native/spark-expr/src/array_funcs/sort_array.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_remove_floating_point.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/sort_array_floating_point.sql diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md index cd167885d05..e0a6eee16e6 100644 --- a/docs/source/contributor-guide/expression-audits/array_funcs.md +++ b/docs/source/contributor-guide/expression-audits/array_funcs.md @@ -122,9 +122,10 @@ ## 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` calls `array_remove_all`, which Comet's `SparkArrayRemove` replaces. For elements with a `FLOAT` or `DOUBLE` at any depth it 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. ## array_repeat @@ -196,8 +197,9 @@ ## 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 takes Spark's `containsNull` from the serde and sorts as Spark's generated code does: a stable sort in Spark's SQL ordering, except that an ascending sort of `FLOAT` or `DOUBLE` elements that cannot be null follows `java.util.Arrays.sort`, which puts `-0.0` before `0.0`. This holds in strict floating-point mode too. Other supported element types use DataFusion's `array_sort`, and unsupported ones route through the codegen dispatcher. `ascendingOrder` may be any foldable boolean, which the serde evaluates. [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 57a9ffe8a83..7ed704bb54d 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 1487e4634e0..45ed0b6807c 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 38fb80f22ec..f36546c2377 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 38f10208bb8..a97334b56f1 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 00000000000..b9468e1dd1a --- /dev/null +++ b/native/spark-expr/benches/float_arrays.rs @@ -0,0 +1,147 @@ +// 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; + +/// `ROWS` lists of `len` doubles, with every tenth list and every seventh element null if +/// `nulls`. Values repeat so that `array_remove` finds something to remove. +fn lists(len: usize, nulls: bool) -> ArrayRef { + let value = |i: usize| ((i * 7919) % 101) as f64 * 0.5 + 1.0; + 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, 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(true))), + ], + &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, scalar(ScalarValue::Float64(Some(26.0)))], + &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, + ), + ]; + 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 00000000000..d2dee0ed216 --- /dev/null +++ b/native/spark-expr/src/array_funcs/array_remove.rs @@ -0,0 +1,522 @@ +// 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 crate::float_semantics::{compare_floats, has_float_leaf, 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, FieldRef, Float32Type, Float64Type}; +use datafusion::common::{exec_err, Result, ScalarValue}; +use datafusion::functions_nested::remove::array_remove_all_udf; +use datafusion::logical_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, +}; +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` compares +/// the bits instead, so it keeps a `-0.0` when removing `0.0`, and a NaN whose bits differ. Other +/// element types go to DataFusion's implementation, which this replaces in Comet's registry. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkArrayRemove { + datafusion_udf: Arc, +} + +impl Default for SparkArrayRemove { + fn default() -> Self { + Self { + datafusion_udf: array_remove_all_udf(), + } + } +} + +impl ScalarUDFImpl for SparkArrayRemove { + fn name(&self) -> &str { + "array_remove_all" + } + + fn signature(&self) -> &Signature { + self.datafusion_udf.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + + /// DataFusion's version answers only this, with the array's field made nullable when either + /// argument is. + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + self.datafusion_udf.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [array, value] = args.args.as_slice() else { + return exec_err!("array_remove takes exactly two arguments"); + }; + // Spark arrays use Arrow's 32-bit List layout. + let DataType::List(field) = array.data_type() else { + return self.datafusion_udf.invoke_with_args(args); + }; + if !has_float_leaf(field.data_type()) { + return self.datafusion_udf.invoke_with_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)?; + let result = match value { + ColumnarValue::Scalar(needle) => remove_constant(array.as_list::(), needle)?, + ColumnarValue::Array(value) => array_remove(array.as_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 { + match array.value_type() { + DataType::Float32 => Ok(remove_floats::(array, value)), + DataType::Float64 => Ok(remove_floats::(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 { + match needle { + ScalarValue::Float32(Some(needle)) => remove_float_constant::(array, *needle), + ScalarValue::Float64(Some(needle)) => remove_float_constant::(array, *needle), + _ => array_remove(array, &needle.to_array_of_size(array.len())?), + } +} + +fn remove_float_constant(array: &ListArray, needle: T::Native) -> Result +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::(); + 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 mut kept_offsets = Vec::with_capacity(array.len() + 1); + kept_offsets.push(0i32); + let mut kept = 0i32; + for bounds in offsets.windows(2) { + let start = bounds[0] as usize - first; + let length = bounds[1] as usize - first - start; + kept += keep + .inner() + .count_set_bits_offset(keep.offset() + start, length) as i32; + kept_offsets.push(kept); + } + let kept_values = filter(&values, &BooleanArray::new(keep, None))?; + Ok(with_values( + array, + kept_offsets, + kept_values, + array.nulls().cloned(), + )) +} + +/// 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(array: &ListArray, value: &ArrayRef) -> ArrayRef +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let values = array.values().as_primitive::(); + let needles = value.as_primitive::().values(); + let (buffer, nulls) = (values.values(), values.nulls()); + let row_nulls = NullBuffer::union(array.nulls(), value.nulls()); + let mut kept: Vec = 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::::new( + kept.into(), + validity.map(|mut validity| NullBuffer::new(validity.finish())), + ); + with_values(array, offsets, 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(array: &ListArray, value: &ArrayRef, equal: F) -> Result +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, offsets, kept_values, row_nulls)) +} + +/// A list with the type of `array` and the given nulls, holding `values` at `offsets`. +fn with_values( + array: &ListArray, + offsets: Vec, + values: ArrayRef, + nulls: Option, +) -> ArrayRef { + let DataType::List(field) = array.data_type() else { + unreachable!("array_remove takes a List"); + }; + Arc::new(ListArray::new( + Arc::clone(field), + OffsetBuffer::new(offsets.into()), + values, + nulls, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Float32Array, Float64Array, Int32Array, StructArray}; + use arrow::datatypes::{Field, Fields}; + use datafusion::config::ConfigOptions; + + /// A NaN with the sign bit set, which arithmetic produces on x86-64. + const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); + + const EDGE_VALUES: [Option; 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), + // A NaN with a payload. + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + None, + ]; + + fn invoke(array: ColumnarValue, value: ColumnarValue, rows: usize) -> Result { + 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()), + }) + } + + fn list(rows: &[Option>>]) -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::( + rows.iter().cloned(), + )) + } + + fn bits(array: &ArrayRef) -> Vec>>> { + array + .as_list::() + .iter() + .map(|row| { + row.map(|row| { + row.as_primitive::() + .iter() + .map(|v| v.map(f64::to_bits)) + .collect() + }) + }) + .collect() + } + + /// Spark's result for one row: drop the non-null elements that `genEqual` the value. + fn spark_remove(row: &[Option], value: f64) -> Vec> { + 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> = EDGE_VALUES.to_vec(); + let removable: Vec = EDGE_VALUES.iter().flatten().copied().collect(); + let rows: Vec>>> = + 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>>> = 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>>> = 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(()) + } + + #[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::(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| -> 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 row = result.as_list::().value(0); + let x = row + .as_struct() + .column(0) + .as_primitive::() + .clone(); + let actual: Vec = x.values().iter().map(|v| v.to_bits()).collect(); + let expected: Vec = expected.iter().map(|v| v.to_bits()).collect(); + assert_eq!(actual, expected, "remove {value}"); + } + Ok(()) + } + + /// The planner asks the UDF for its return field, which DataFusion's version only answers + /// through `return_field_from_args`. + #[test] + fn return_field() -> Result<()> { + let udf = ScalarUDF::new_from_impl(SparkArrayRemove::default()); + let list = Arc::new(Field::new( + "array", + DataType::new_list(DataType::Float64, true), + false, + )); + let value = Arc::new(Field::new("value", DataType::Float64, true)); + let field = udf.return_field_from_args(ReturnFieldArgs { + arg_fields: &[Arc::clone(&list), value], + scalar_arguments: &[None, None], + })?; + assert_eq!(field.data_type(), list.data_type()); + assert!(field.is_nullable(), "a null value gives a null array"); + Ok(()) + } + + #[test] + fn other_types_go_to_datafusion() -> Result<()> { + let array = Arc::new(ListArray::from_iter_primitive::< + arrow::datatypes::Int32Type, + _, + _, + >(vec![Some(vec![Some(1), Some(2), Some(1)])])); + let result = invoke( + ColumnarValue::Array(array), + ColumnarValue::Array(Arc::new(Int32Array::from(vec![1]))), + 1, + )? + .into_array(1)?; + let row = result.as_list::().value(0); + assert_eq!( + row.as_primitive::() + .values() + .as_ref(), + &[2] + ); + Ok(()) + } +} diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index f5d87f09a8c..371d6bae52c 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,4 @@ 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; 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 00000000000..6090d6c465b --- /dev/null +++ b/native/spark-expr/src/array_funcs/sort_array.rs @@ -0,0 +1,560 @@ +// 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 crate::float_semantics::{compare_floats, 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, 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. Spark's generated code makes +/// one exception: it sorts an ascending array of `FLOAT` or `DOUBLE` that cannot hold a null with +/// `java.util.Arrays.sort`, which puts `-0.0` before `0.0`. 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 Spark's `containsNull` for the +/// array, 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 { + Ok(arg_types[0].clone()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [array, ascending, contains_null] = args.args.as_slice() else { + return exec_err!("spark_sort_array takes exactly three arguments"); + }; + let (Some(ascending), Some(contains_null)) = (constant(ascending), constant(contains_null)) + else { + return exec_err!("spark_sort_array takes constant ascending and containsNull flags"); + }; + let is_scalar = matches!(array, ColumnarValue::Scalar(_)); + let rows = if is_scalar { 1 } else { args.number_rows }; + // Spark arrays use Arrow's 32-bit List layout. + let result = sort_array( + array.to_array(rows)?.as_list::(), + ascending, + contains_null, + )?; + if is_scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(result)) + } + } +} + +fn constant(value: &ColumnarValue) -> Option { + match value { + ColumnarValue::Scalar(ScalarValue::Boolean(Some(value))) => Some(*value), + _ => None, + } +} + +fn sort_array(array: &ListArray, ascending: bool, contains_null: bool) -> Result { + let java_order = ascending && !contains_null; + match array.value_type() { + DataType::Float32 => Ok(sort_floats::(array, ascending, java_order)), + DataType::Float64 => Ok(sort_floats::(array, ascending, java_order)), + _ => { + let values = array.values(); + let compare = spark_comparator(values.as_ref(), values.as_ref())?; + sort_by(array, ascending, compare) + } + } +} + +/// The order of `Double.compare`, which `java.util.Arrays.sort` uses: Spark's SQL ordering, except +/// that `-0.0` comes before `0.0`. +fn java_order(a: T, b: T) -> Ordering { + match compare_floats(a, b) { + // Equal values that are not NaN differ in sign only when they are -0.0 and 0.0. + Ordering::Equal if !a.is_nan() => b.is_sign_negative().cmp(&a.is_sign_negative()), + ordering => ordering, + } +} + +/// 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(array: &ListArray, ascending: bool, java: bool) -> ArrayRef +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let values = array.values().as_primitive::(); + 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.iter().map(|offset| offset - offsets[0]).collect(); + 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 row = &mut sorted[start - base..end - base]; + if !sort_row(row, ascending, java) { + restore_ties(row, buffer[start..end].iter().copied(), ascending, java); + } + } + } + let values = PrimitiveArray::::new(sorted.into(), None); + return with_values(array, rebased, Arc::new(values)); + }; + // Each row's valid values are copied after its nulls when ascending, or before them when + // descending, and sorted there. The copy writes every element and moves past only the valid + // ones, so it can write one place past a row's valid values, into the spare element at the + // end. A null row's elements stay null, hidden by the list's null. + let mut sorted = vec![T::Native::default(); total + 1]; + 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 valid = nulls + .buffer() + .count_set_bits_offset(nulls.offset() + start, end - start); + let leading_nulls = if ascending { end - start - valid } else { 0 }; + let first = start - base + leading_nulls; + let mut next = first; + for index in start..end { + sorted[next] = buffer[index]; + next += usize::from(nulls.is_valid(index)); + } + let row = &mut sorted[first..next]; + if !sort_row(row, ascending, java) { + let original = (start..end) + .filter(|&index| nulls.is_valid(index)) + .map(|index| buffer[index]); + restore_ties(row, original, ascending, java); + } + validity.append_n(leading_nulls, false); + validity.append_n(valid, true); + validity.append_n(end - start - valid - leading_nulls, false); + } + sorted.truncate(total); + let values = PrimitiveArray::::new(sorted.into(), Some(NullBuffer::new(validity.finish()))); + with_values(array, rebased, Arc::new(values)) +} + +/// 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 in Spark's SQL ordering, reversed when descending, or in [`java_order`], +/// and returns whether the sort was stable, as Spark's is. A longer row sorts faster with an +/// unstable sort, after which [`restore_ties`] gives the stable result. +fn sort_row(row: &mut [N], ascending: bool, java: bool) -> bool { + let stable = row.len() <= STABLE_SORT_MAX_LEN; + match (java, ascending, stable) { + (true, _, true) => row.sort_by(|a, b| java_order(*a, *b)), + (true, _, false) => row.sort_unstable_by(|a, b| java_order(*a, *b)), + (false, true, true) => row.sort_by(|a, b| compare_floats(*a, *b)), + (false, true, false) => row.sort_unstable_by(|a, b| compare_floats(*a, *b)), + (false, false, true) => row.sort_by(|a, b| compare_floats(*b, *a)), + (false, false, false) => row.sort_unstable_by(|a, b| compare_floats(*b, *a)), + } + stable +} + +/// Puts the tied elements of a row sorted by [`sort_row`] 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 when ascending and first +/// when descending, and zeros of either sign, except in [`java_order`]. Each forms one run. +fn restore_ties( + row: &mut [N], + original: impl Iterator + Clone, + ascending: bool, + java: bool, +) { + let len = row.len(); + let nans = if ascending { + len - row.iter().rev().take_while(|v| v.is_nan()).count()..len + } else { + 0..row.iter().take_while(|v| v.is_nan()).count() + }; + restore_order(&mut row[nans], original.clone().filter(|v| v.is_nan())); + if !java { + let zero = N::zero(); + let first = if ascending { + row.partition_point(|v| *v < zero) + } else { + row.partition_point(|v| v.is_nan() || *v > zero) + }; + 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(run: &mut [N], values: impl Iterator) { + if run.len() > 1 { + run.iter_mut() + .zip(values) + .for_each(|(slot, value)| *slot = value); + } +} + +/// A list with the type and nulls of `array`, holding `values` at `offsets`. +fn with_values(array: &ListArray, offsets: Vec, values: ArrayRef) -> ArrayRef { + let DataType::List(field) = array.data_type() else { + unreachable!("sort_array takes a List"); + }; + Arc::new(ListArray::new( + Arc::clone(field), + OffsetBuffer::new(offsets.into()), + values, + array.nulls().cloned(), + )) +} + +/// Sorts the elements of each row stably with `compare`, reversed when descending, after putting +/// null elements first when ascending and last when descending, as Spark's comparators do. +fn sort_by(array: &ListArray, ascending: bool, compare: F) -> Result +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 nulls = values.logical_nulls(); + let is_null = |index: usize| nulls.as_ref().is_some_and(|nulls| nulls.is_null(index)); + let mut indices: Vec = 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); + match (is_null(a), is_null(b)) { + (true, true) => Ordering::Equal, + (true, false) if ascending => Ordering::Less, + (true, false) => Ordering::Greater, + (false, true) if ascending => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) if ascending => compare(a, b), + (false, false) => compare(b, a), + } + }); + } + offsets.push(indices.len() as i32); + } + let sorted = take(values.as_ref(), &UInt32Array::from(indices), None)?; + Ok(with_values(array, offsets, sorted)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Float32Array, Float64Array, StructArray}; + use arrow::datatypes::{Field, Fields}; + use datafusion::config::ConfigOptions; + + /// A NaN with the sign bit set, which arithmetic produces on x86-64. + const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); + + const EDGE_VALUES: [Option; 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), + // A NaN with a payload. + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + None, + ]; + + fn invoke(array: ArrayRef, ascending: bool, contains_null: bool) -> Result { + 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(contains_null))), + ], + arg_fields: vec![], + number_rows: rows, + return_field, + config_options: Arc::new(ConfigOptions::default()), + })? + .into_array(rows) + } + + fn list(rows: &[Option>>]) -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::( + rows.iter().cloned(), + )) + } + + fn bits(array: &ArrayRef) -> Vec>>> { + array + .as_list::() + .iter() + .map(|row| { + row.map(|row| { + row.as_primitive::() + .iter() + .map(|v| v.map(f64::to_bits)) + .collect() + }) + }) + .collect() + } + + fn row_bits(row: &[Option]) -> Vec> { + 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>>> { + 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>>> { + 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], ascending: bool) -> Vec> { + 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>>> = rows + .iter() + .map(|row| row.as_ref().map(|row| spark_sort(row, ascending))) + .collect(); + let result = invoke(list(&rows), ascending, true)?; + assert_eq!(bits(&result), expected, "ascending={ascending}"); + let result = invoke(list(&rows).slice(37, 120), ascending, true)?; + 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 non_null_ascending_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>>> = rows + .iter() + .map(|row| { + row.as_ref().map(|row| { + let mut sorted: Vec = 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, false)?; + assert_eq!(bits(&result), expected); + Ok(()) + } + + /// Spark's generated code sorts an ascending array that cannot hold a null with + /// `java.util.Arrays.sort`, which puts -0.0 before 0.0. Descending still ties them. + #[test] + fn non_null_ascending_puts_negative_zero_first() -> Result<()> { + let rows = vec![Some(vec![ + Some(0.0), + Some(1.0), + Some(-0.0), + Some(NEGATIVE_NAN), + ])]; + let cases = [ + (true, false, vec![-0.0, 0.0, 1.0, NEGATIVE_NAN]), + (true, true, vec![0.0, -0.0, 1.0, NEGATIVE_NAN]), + (false, false, vec![NEGATIVE_NAN, 1.0, 0.0, -0.0]), + ]; + for (ascending, contains_null, expected) in cases { + let result = invoke(list(&rows), ascending, contains_null)?; + let expected: Vec> = expected.iter().map(|v| Some(v.to_bits())).collect(); + assert_eq!( + bits(&result), + vec![Some(expected)], + "ascending={ascending} contains_null={contains_null}" + ); + } + Ok(()) + } + + #[test] + fn float32() -> Result<()> { + let array = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![ + Some(f32::NAN), + Some(0.0f32), + Some(f32::from_bits(0xffc0_0000)), + Some(-0.0), + None, + ]), + ])); + let result = invoke(array, true, true)?; + let row = result.as_list::().value(0); + let row: &Float32Array = row.as_primitive(); + let actual: Vec> = 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> = 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]), + ] { + // Not primitive, so Spark sorts in its SQL ordering whatever containsNull says. + let result = invoke(Arc::clone(&array), ascending, false)?; + let row = result.as_list::().value(0); + let x = row + .as_struct() + .column(0) + .as_primitive::() + .clone(); + let actual: Vec = x.values().iter().map(|v| v.to_bits()).collect(); + let expected: Vec = expected.iter().map(|v| v.to_bits()).collect(); + assert_eq!(actual, 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 6ccaff2b2dc..7b3c9806399 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::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/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a570b033258..63a9b09b35a 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1079,10 +1079,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 8cdabebb85d..6a07a5a36f0 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._ @@ -140,10 +139,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 +158,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 +176,28 @@ 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, whose generated code also sorts differently when the elements cannot be null, so it + // takes containsNull, which the native array type does not carry. + scalarFunctionExprToProto( + "spark_sort_array", + arrayExprProto, + exprToProtoInternal(Literal(ascending), inputs, binding), + exprToProtoInternal(Literal(arrayType.containsNull), 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 00000000000..aa7bbb41bd5 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_remove_floating_point.sql @@ -0,0 +1,59 @@ +-- 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. + +-- ConfigMatrix: spark.comet.exec.strictFloatingPoint=false,true + +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 + +-- 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 ece58525d46..a6e9499c047 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 03d4d3120a6..1e688857a9e 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 00000000000..76924b35b65 --- /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. + +-- Strict floating-point mode no longer needs the codegen dispatcher for these. +-- ConfigMatrix: spark.comet.exec.strictFloatingPoint=false,true +-- Config: spark.comet.expression.SortArray.allowIncompatible=false + +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 + +-- 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 db3453bd235..2b0a3aab2bd 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))) From da3410aad5a4493bc9045187a87b9974975c28aa Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 1 Oct 2026 12:17:42 -0600 Subject: [PATCH 2/4] 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. --- native/spark-expr/benches/float_arrays.rs | 24 +++++++++++++-- .../src/array_funcs/array_remove.rs | 30 +++++++++++++++---- .../spark-expr/src/array_funcs/sort_array.rs | 15 ++++++---- .../array/array_remove_floating_point.sql | 6 ++++ .../array/sort_array_floating_point.sql | 4 +++ 5 files changed, 65 insertions(+), 14 deletions(-) diff --git a/native/spark-expr/benches/float_arrays.rs b/native/spark-expr/benches/float_arrays.rs index b9468e1dd1a..6cdc94ba3c0 100644 --- a/native/spark-expr/benches/float_arrays.rs +++ b/native/spark-expr/benches/float_arrays.rs @@ -35,10 +35,14 @@ 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`. Values repeat so that `array_remove` finds something to remove. +/// `nulls`. fn lists(len: usize, nulls: bool) -> ArrayRef { - let value = |i: usize| ((i * 7919) % 101) as f64 * 0.5 + 1.0; let values = if nulls { Float64Array::from_iter((0..ROWS * len).map(|i| (i % 7 != 0).then(|| value(i)))) } else { @@ -96,7 +100,14 @@ fn criterion_benchmark(c: &mut Criterion) { &list, ); let remove_args = args( - vec![column, scalar(ScalarValue::Float64(Some(26.0)))], + 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 = [ @@ -114,6 +125,13 @@ fn criterion_benchmark(c: &mut Criterion) { &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| { diff --git a/native/spark-expr/src/array_funcs/array_remove.rs b/native/spark-expr/src/array_funcs/array_remove.rs index d2dee0ed216..f93228ab7e6 100644 --- a/native/spark-expr/src/array_funcs/array_remove.rs +++ b/native/spark-expr/src/array_funcs/array_remove.rs @@ -102,9 +102,10 @@ impl ScalarUDFImpl for SparkArrayRemove { } fn array_remove(array: &ListArray, value: &ArrayRef) -> Result { - match array.value_type() { - DataType::Float32 => Ok(remove_floats::(array, value)), - DataType::Float64 => Ok(remove_floats::(array, value)), + // 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::(array, value)), + (DataType::Float64, DataType::Float64) => Ok(remove_floats::(array, value)), _ => { let equal = spark_equality(array.values().as_ref(), value.as_ref())?; remove_where(array, value, equal) @@ -115,9 +116,13 @@ fn array_remove(array: &ListArray, value: &ArrayRef) -> Result { /// [`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 { - match needle { - ScalarValue::Float32(Some(needle)) => remove_float_constant::(array, *needle), - ScalarValue::Float64(Some(needle)) => remove_float_constant::(array, *needle), + match (array.value_type(), needle) { + (DataType::Float32, ScalarValue::Float32(Some(needle))) => { + remove_float_constant::(array, *needle) + } + (DataType::Float64, ScalarValue::Float64(Some(needle))) => { + remove_float_constant::(array, *needle) + } _ => array_remove(array, &needle.to_array_of_size(array.len())?), } } @@ -497,6 +502,19 @@ mod tests { 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()); + } + } + #[test] fn other_types_go_to_datafusion() -> Result<()> { let array = Arc::new(ListArray::from_iter_primitive::< diff --git a/native/spark-expr/src/array_funcs/sort_array.rs b/native/spark-expr/src/array_funcs/sort_array.rs index 6090d6c465b..388f0facc81 100644 --- a/native/spark-expr/src/array_funcs/sort_array.rs +++ b/native/spark-expr/src/array_funcs/sort_array.rs @@ -78,12 +78,12 @@ impl ScalarUDFImpl for SparkSortArray { }; 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 result = sort_array( - array.to_array(rows)?.as_list::(), - ascending, - contains_null, - )?; + let Some(list) = array.as_list_opt::() else { + return exec_err!("spark_sort_array takes a list, got {}", array.data_type()); + }; + let result = sort_array(list, ascending, contains_null)?; if is_scalar { Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( &result, 0, @@ -496,6 +496,11 @@ mod tests { Ok(()) } + #[test] + fn non_list_is_an_error() { + assert!(invoke(Arc::new(Float64Array::from(vec![1.0])), true, true).is_err()); + } + #[test] fn float32() -> Result<()> { let array = Arc::new(ListArray::from_iter_primitive::(vec![ 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 index aa7bbb41bd5..b7443ab0603 100644 --- 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 @@ -45,6 +45,12 @@ 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 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 index 76924b35b65..39439d25214 100644 --- 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 @@ -50,6 +50,10 @@ 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)), From 98d3270d3cd645dbab1aa55fa302d0611be4ab05 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 1 Oct 2026 13:03:59 -0600 Subject: [PATCH 3/4] 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 | 6 +- native/spark-expr/benches/float_arrays.rs | 2 +- .../src/array_funcs/array_remove.rs | 241 +++++-------- native/spark-expr/src/array_funcs/mod.rs | 68 ++++ .../spark-expr/src/array_funcs/sort_array.rs | 336 ++++++++---------- native/spark-expr/src/float_semantics/mod.rs | 18 +- .../scala/org/apache/comet/serde/arrays.scala | 20 +- .../array/array_remove_floating_point.sql | 2 - .../array/sort_array_floating_point.sql | 4 - 9 files changed, 342 insertions(+), 355 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md index e0a6eee16e6..0911c45aff0 100644 --- a/docs/source/contributor-guide/expression-audits/array_funcs.md +++ b/docs/source/contributor-guide/expression-audits/array_funcs.md @@ -125,7 +125,8 @@ - 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` calls `array_remove_all`, which Comet's `SparkArrayRemove` replaces. For elements with a `FLOAT` or `DOUBLE` at any depth it 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. +- 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_arrays.rs`. ## array_repeat @@ -200,6 +201,7 @@ - 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 takes Spark's `containsNull` from the serde and sorts as Spark's generated code does: a stable sort in Spark's SQL ordering, except that an ascending sort of `FLOAT` or `DOUBLE` elements that cannot be null follows `java.util.Arrays.sort`, which puts `-0.0` before `0.0`. This holds in strict floating-point mode too. Other supported element types use DataFusion's `array_sort`, and unsupported ones route through the codegen dispatcher. `ascendingOrder` may be any foldable boolean, which the serde evaluates. +- 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 the codegen dispatcher. `ascendingOrder` may be any foldable boolean, which the serde evaluates. +- 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 time than DataFusion's `array_sort`. Benchmark: `benches/float_arrays.rs`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/native/spark-expr/benches/float_arrays.rs b/native/spark-expr/benches/float_arrays.rs index 6cdc94ba3c0..44a4d78a2fc 100644 --- a/native/spark-expr/benches/float_arrays.rs +++ b/native/spark-expr/benches/float_arrays.rs @@ -87,7 +87,7 @@ fn criterion_benchmark(c: &mut Criterion) { vec![ column.clone(), scalar(ScalarValue::Boolean(Some(true))), - scalar(ScalarValue::Boolean(Some(true))), + scalar(ScalarValue::Boolean(Some(false))), ], &list, ); diff --git a/native/spark-expr/src/array_funcs/array_remove.rs b/native/spark-expr/src/array_funcs/array_remove.rs index f93228ab7e6..aac6dd4c8cd 100644 --- a/native/spark-expr/src/array_funcs/array_remove.rs +++ b/native/spark-expr/src/array_funcs/array_remove.rs @@ -17,17 +17,17 @@ use std::sync::Arc; -use crate::float_semantics::{compare_floats, has_float_leaf, spark_equality}; +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, FieldRef, Float32Type, Float64Type}; -use datafusion::common::{exec_err, Result, ScalarValue}; -use datafusion::functions_nested::remove::array_remove_all_udf; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; +use datafusion::common::{exec_err, utils::take_function_args, Result, ScalarValue}; use datafusion::logical_expr::{ - ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; use num::Float; @@ -35,61 +35,50 @@ use num::Float; /// /// 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` compares -/// the bits instead, so it keeps a `-0.0` when removing `0.0`, and a NaN whose bits differ. Other -/// element types go to DataFusion's implementation, which this replaces in Comet's registry. +/// 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 { - datafusion_udf: Arc, + signature: Signature, } impl Default for SparkArrayRemove { fn default() -> Self { Self { - datafusion_udf: array_remove_all_udf(), + signature: Signature::any(2, Volatility::Immutable), } } } impl ScalarUDFImpl for SparkArrayRemove { fn name(&self) -> &str { - "array_remove_all" + "spark_array_remove" } fn signature(&self) -> &Signature { - self.datafusion_udf.signature() + &self.signature } fn return_type(&self, arg_types: &[DataType]) -> Result { Ok(arg_types[0].clone()) } - /// DataFusion's version answers only this, with the array's field made nullable when either - /// argument is. - fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { - self.datafusion_udf.return_field_from_args(args) - } - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let [array, value] = args.args.as_slice() else { - return exec_err!("array_remove takes exactly two arguments"); - }; - // Spark arrays use Arrow's 32-bit List layout. - let DataType::List(field) = array.data_type() else { - return self.datafusion_udf.invoke_with_args(args); - }; - if !has_float_leaf(field.data_type()) { - return self.datafusion_udf.invoke_with_args(args); - } + 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::() else { + return exec_err!("spark_array_remove takes a list, got {}", array.data_type()); + }; let result = match value { - ColumnarValue::Scalar(needle) => remove_constant(array.as_list::(), needle)?, - ColumnarValue::Array(value) => array_remove(array.as_list::(), 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( @@ -152,17 +141,7 @@ where None => differs, }; // A null row keeps its elements in the values; the list's null hides them. - let mut kept_offsets = Vec::with_capacity(array.len() + 1); - kept_offsets.push(0i32); - let mut kept = 0i32; - for bounds in offsets.windows(2) { - let start = bounds[0] as usize - first; - let length = bounds[1] as usize - first - start; - kept += keep - .inner() - .count_set_bits_offset(keep.offset() + start, length) as i32; - kept_offsets.push(kept); - } + let kept_offsets = kept_offsets(&keep, offsets); let kept_values = filter(&values, &BooleanArray::new(keep, None))?; Ok(with_values( array, @@ -172,6 +151,25 @@ where )) } +/// 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) -> OffsetBuffer { + 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::>().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(array: &ListArray, value: &ArrayRef) -> ArrayRef @@ -215,7 +213,12 @@ where kept.into(), validity.map(|mut validity| NullBuffer::new(validity.finish())), ); - with_values(array, offsets, Arc::new(kept), row_nulls) + 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 @@ -247,51 +250,23 @@ where offsets.push(kept); } let kept_values = filter(values, &BooleanArray::new(keep.finish(), None))?; - Ok(with_values(array, offsets, kept_values, row_nulls)) -} - -/// A list with the type of `array` and the given nulls, holding `values` at `offsets`. -fn with_values( - array: &ListArray, - offsets: Vec, - values: ArrayRef, - nulls: Option, -) -> ArrayRef { - let DataType::List(field) = array.data_type() else { - unreachable!("array_remove takes a List"); - }; - Arc::new(ListArray::new( - Arc::clone(field), + Ok(with_values( + array, OffsetBuffer::new(offsets.into()), - values, - nulls, + kept_values, + row_nulls, )) } #[cfg(test)] mod tests { + use super::super::test_util::{bits, first_row_field_bits, list, EDGE_VALUES}; use super::*; - use arrow::array::{Float32Array, Float64Array, Int32Array, StructArray}; + use crate::float_semantics::NEGATIVE_NAN; + use arrow::array::{Float32Array, Float64Array, StructArray}; use arrow::datatypes::{Field, Fields}; use datafusion::config::ConfigOptions; - /// A NaN with the sign bit set, which arithmetic produces on x86-64. - const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); - - const EDGE_VALUES: [Option; 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), - // A NaN with a payload. - Some(f64::from_bits(0x7ff0_0000_0000_0001)), - None, - ]; - fn invoke(array: ColumnarValue, value: ColumnarValue, rows: usize) -> Result { let return_field = Arc::new(Field::new("result", array.data_type(), true)); SparkArrayRemove::default().invoke_with_args(ScalarFunctionArgs { @@ -303,27 +278,6 @@ mod tests { }) } - fn list(rows: &[Option>>]) -> ArrayRef { - Arc::new(ListArray::from_iter_primitive::( - rows.iter().cloned(), - )) - } - - fn bits(array: &ArrayRef) -> Vec>>> { - array - .as_list::() - .iter() - .map(|row| { - row.map(|row| { - row.as_primitive::() - .iter() - .map(|v| v.map(f64::to_bits)) - .collect() - }) - }) - .collect() - } - /// Spark's result for one row: drop the non-null elements that `genEqual` the value. fn spark_remove(row: &[Option], value: f64) -> Vec> { row.iter() @@ -389,6 +343,45 @@ mod tests { 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>>> = 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>>> = 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![ @@ -469,39 +462,12 @@ mod tests { 1, )? .into_array(1)?; - let row = result.as_list::().value(0); - let x = row - .as_struct() - .column(0) - .as_primitive::() - .clone(); - let actual: Vec = x.values().iter().map(|v| v.to_bits()).collect(); let expected: Vec = expected.iter().map(|v| v.to_bits()).collect(); - assert_eq!(actual, expected, "remove {value}"); + assert_eq!(first_row_field_bits(&result), expected, "remove {value}"); } Ok(()) } - /// The planner asks the UDF for its return field, which DataFusion's version only answers - /// through `return_field_from_args`. - #[test] - fn return_field() -> Result<()> { - let udf = ScalarUDF::new_from_impl(SparkArrayRemove::default()); - let list = Arc::new(Field::new( - "array", - DataType::new_list(DataType::Float64, true), - false, - )); - let value = Arc::new(Field::new("value", DataType::Float64, true)); - let field = udf.return_field_from_args(ReturnFieldArgs { - arg_fields: &[Arc::clone(&list), value], - scalar_arguments: &[None, None], - })?; - assert_eq!(field.data_type(), list.data_type()); - assert!(field.is_nullable(), "a null value gives a null array"); - 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] @@ -514,27 +480,4 @@ mod tests { assert!(invoke(ColumnarValue::Array(Arc::clone(&array)), value, 1).is_err()); } } - - #[test] - fn other_types_go_to_datafusion() -> Result<()> { - let array = Arc::new(ListArray::from_iter_primitive::< - arrow::datatypes::Int32Type, - _, - _, - >(vec![Some(vec![Some(1), Some(2), Some(1)])])); - let result = invoke( - ColumnarValue::Array(array), - ColumnarValue::Array(Arc::new(Int32Array::from(vec![1]))), - 1, - )? - .into_array(1)?; - let row = result.as_list::().value(0); - assert_eq!( - row.as_primitive::() - .values() - .as_ref(), - &[2] - ); - Ok(()) - } } diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index 371d6bae52c..ab400790368 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -46,3 +46,71 @@ 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, + values: ArrayRef, + nulls: Option, +) -> 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; 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>>]) -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::( + 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>>> { + array + .as_list::() + .iter() + .map(|row| { + row.map(|row| { + row.as_primitive::() + .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 { + let row = array.as_list::().value(0); + let field = row.as_struct().column(0).as_primitive::(); + 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 index 388f0facc81..a1b55fa5f7c 100644 --- a/native/spark-expr/src/array_funcs/sort_array.rs +++ b/native/spark-expr/src/array_funcs/sort_array.rs @@ -18,14 +18,17 @@ use std::cmp::Ordering; use std::sync::Arc; -use crate::float_semantics::{compare_floats, spark_comparator}; +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, Result, ScalarValue}; +use datafusion::common::{exec_err, utils::take_function_args, Result, ScalarValue}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; @@ -35,13 +38,13 @@ use num::Float; /// /// 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. Spark's generated code makes -/// one exception: it sorts an ascending array of `FLOAT` or `DOUBLE` that cannot hold a null with -/// `java.util.Arrays.sort`, which puts `-0.0` before `0.0`. DataFusion's `array_sort` uses IEEE -/// 754 total order instead, in which a NaN with the sign bit set sorts first. +/// 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 Spark's `containsNull` for the -/// array, which Arrow's field nullability does not carry. +/// 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, @@ -69,12 +72,17 @@ impl ScalarUDFImpl for SparkSortArray { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let [array, ascending, contains_null] = args.args.as_slice() else { - return exec_err!("spark_sort_array takes exactly three arguments"); - }; - let (Some(ascending), Some(contains_null)) = (constant(ascending), constant(contains_null)) - else { - return exec_err!("spark_sort_array takes constant ascending and containsNull flags"); + 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 }; @@ -83,7 +91,7 @@ impl ScalarUDFImpl for SparkSortArray { let Some(list) = array.as_list_opt::() else { return exec_err!("spark_sort_array takes a list, got {}", array.data_type()); }; - let result = sort_array(list, ascending, contains_null)?; + let result = sort_array(list, order)?; if is_scalar { Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( &result, 0, @@ -101,32 +109,33 @@ fn constant(value: &ColumnarValue) -> Option { } } -fn sort_array(array: &ListArray, ascending: bool, contains_null: bool) -> Result { - let java_order = ascending && !contains_null; +/// 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 { match array.value_type() { - DataType::Float32 => Ok(sort_floats::(array, ascending, java_order)), - DataType::Float64 => Ok(sort_floats::(array, ascending, java_order)), + DataType::Float32 => Ok(sort_floats::(array, order)), + DataType::Float64 => Ok(sort_floats::(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, ascending, compare) + sort_by(array, order != Order::Descending, compare) } } } -/// The order of `Double.compare`, which `java.util.Arrays.sort` uses: Spark's SQL ordering, except -/// that `-0.0` comes before `0.0`. -fn java_order(a: T, b: T) -> Ordering { - match compare_floats(a, b) { - // Equal values that are not NaN differ in sign only when they are -0.0 and 0.0. - Ordering::Equal if !a.is_nan() => b.is_sign_negative().cmp(&a.is_sign_negative()), - ordering => ordering, - } -} - /// 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(array: &ListArray, ascending: bool, java: bool) -> ArrayRef +fn sort_floats(array: &ListArray, order: Order) -> ArrayRef where T: ArrowPrimitiveType, T::Native: Float, @@ -136,27 +145,25 @@ where let offsets = array.offsets(); let base = offsets[0] as usize; let total = offsets[array.len()] as usize - base; - let rebased = offsets.iter().map(|offset| offset - offsets[0]).collect(); + 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 row = &mut sorted[start - base..end - base]; - if !sort_row(row, ascending, java) { - restore_ties(row, buffer[start..end].iter().copied(), ascending, java); - } + let original = buffer[start..end].iter().copied(); + sort_row(&mut sorted[start - base..end - base], original, order); } } let values = PrimitiveArray::::new(sorted.into(), None); - return with_values(array, rebased, Arc::new(values)); + return with_values(array, rebased, Arc::new(values), array.nulls().cloned()); }; - // Each row's valid values are copied after its nulls when ascending, or before them when - // descending, and sorted there. The copy writes every element and moves past only the valid - // ones, so it can write one place past a row's valid values, into the spare element at the - // end. A null row's elements stay null, hidden by the list's null. - let mut sorted = vec![T::Native::default(); total + 1]; + // 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); @@ -164,80 +171,82 @@ where validity.append_n(end - start, false); continue; } - let valid = nulls - .buffer() - .count_set_bits_offset(nulls.offset() + start, end - start); - let leading_nulls = if ascending { end - start - valid } else { 0 }; - let first = start - base + leading_nulls; - let mut next = first; - for index in start..end { - sorted[next] = buffer[index]; - next += usize::from(nulls.is_valid(index)); - } - let row = &mut sorted[first..next]; - if !sort_row(row, ascending, java) { - let original = (start..end) - .filter(|&index| nulls.is_valid(index)) - .map(|index| buffer[index]); - restore_ties(row, original, ascending, java); - } - validity.append_n(leading_nulls, false); - validity.append_n(valid, true); - validity.append_n(end - start - valid - leading_nulls, false); + 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); } - sorted.truncate(total); let values = PrimitiveArray::::new(sorted.into(), Some(NullBuffer::new(validity.finish()))); - with_values(array, rebased, Arc::new(values)) + 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 in Spark's SQL ordering, reversed when descending, or in [`java_order`], -/// and returns whether the sort was stable, as Spark's is. A longer row sorts faster with an -/// unstable sort, after which [`restore_ties`] gives the stable result. -fn sort_row(row: &mut [N], ascending: bool, java: bool) -> bool { +/// 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(row: &mut [N], original: impl Iterator + Clone, order: Order) { let stable = row.len() <= STABLE_SORT_MAX_LEN; - match (java, ascending, stable) { - (true, _, true) => row.sort_by(|a, b| java_order(*a, *b)), - (true, _, false) => row.sort_unstable_by(|a, b| java_order(*a, *b)), - (false, true, true) => row.sort_by(|a, b| compare_floats(*a, *b)), - (false, true, false) => row.sort_unstable_by(|a, b| compare_floats(*a, *b)), - (false, false, true) => row.sort_by(|a, b| compare_floats(*b, *a)), - (false, false, false) => row.sort_unstable_by(|a, b| compare_floats(*b, *a)), + 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(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)), } - stable } -/// Puts the tied elements of a row sorted by [`sort_row`] back in their order in `original`, the -/// row's values in their original order, as Spark's stable sort leaves them. +/// 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 when ascending and first -/// when descending, and zeros of either sign, except in [`java_order`]. Each forms one run. -fn restore_ties( - row: &mut [N], - original: impl Iterator + Clone, - ascending: bool, - java: bool, -) { +/// 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(row: &mut [N], original: impl Iterator + Clone, order: Order) { let len = row.len(); - let nans = if ascending { - len - row.iter().rev().take_while(|v| v.is_nan()).count()..len - } else { + 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())); - if !java { - let zero = N::zero(); - let first = if ascending { - row.partition_point(|v| *v < zero) - } else { - row.partition_point(|v| v.is_nan() || *v > zero) - }; - let end = first + row[first..].iter().take_while(|v| **v == zero).count(); - restore_order(&mut row[first..end], original.filter(|v| *v == zero)); - } + 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. @@ -249,21 +258,9 @@ fn restore_order(run: &mut [N], values: impl Iterator) { } } -/// A list with the type and nulls of `array`, holding `values` at `offsets`. -fn with_values(array: &ListArray, offsets: Vec, values: ArrayRef) -> ArrayRef { - let DataType::List(field) = array.data_type() else { - unreachable!("sort_array takes a List"); - }; - Arc::new(ListArray::new( - Arc::clone(field), - OffsetBuffer::new(offsets.into()), - values, - array.nulls().cloned(), - )) -} - -/// Sorts the elements of each row stably with `compare`, reversed when descending, after putting -/// null elements first when ascending and last when descending, as Spark's comparators do. +/// 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(array: &ListArray, ascending: bool, compare: F) -> Result where F: Fn(usize, usize) -> Ordering, @@ -272,8 +269,6 @@ where if values.len() > u32::MAX as usize { return exec_err!("sort_array cannot sort more than {} elements", u32::MAX); } - let nulls = values.logical_nulls(); - let is_null = |index: usize| nulls.as_ref().is_some_and(|nulls| nulls.is_null(index)); let mut indices: Vec = Vec::with_capacity(values.len()); let mut offsets = Vec::with_capacity(array.len() + 1); offsets.push(0i32); @@ -283,48 +278,34 @@ where 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); - match (is_null(a), is_null(b)) { - (true, true) => Ordering::Equal, - (true, false) if ascending => Ordering::Less, - (true, false) => Ordering::Greater, - (false, true) if ascending => Ordering::Greater, - (false, true) => Ordering::Less, - (false, false) if ascending => compare(a, b), - (false, false) => compare(b, a), + 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, offsets, sorted)) + 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; - /// A NaN with the sign bit set, which arithmetic produces on x86-64. - const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); - - const EDGE_VALUES: [Option; 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), - // A NaN with a payload. - Some(f64::from_bits(0x7ff0_0000_0000_0001)), - None, - ]; - - fn invoke(array: ArrayRef, ascending: bool, contains_null: bool) -> Result { + fn invoke(array: ArrayRef, ascending: bool, negative_zero_first: bool) -> Result { let rows = array.len(); let return_field = Arc::new(Field::new("result", array.data_type().clone(), true)); SparkSortArray::default() @@ -332,7 +313,7 @@ mod tests { args: vec![ ColumnarValue::Array(array), ColumnarValue::Scalar(ScalarValue::Boolean(Some(ascending))), - ColumnarValue::Scalar(ScalarValue::Boolean(Some(contains_null))), + ColumnarValue::Scalar(ScalarValue::Boolean(Some(negative_zero_first))), ], arg_fields: vec![], number_rows: rows, @@ -342,27 +323,6 @@ mod tests { .into_array(rows) } - fn list(rows: &[Option>>]) -> ArrayRef { - Arc::new(ListArray::from_iter_primitive::( - rows.iter().cloned(), - )) - } - - fn bits(array: &ArrayRef) -> Vec>>> { - array - .as_list::() - .iter() - .map(|row| { - row.map(|row| { - row.as_primitive::() - .iter() - .map(|v| v.map(f64::to_bits)) - .collect() - }) - }) - .collect() - } - fn row_bits(row: &[Option]) -> Vec> { row.iter().map(|v| v.map(f64::to_bits)).collect() } @@ -428,9 +388,9 @@ mod tests { .iter() .map(|row| row.as_ref().map(|row| spark_sort(row, ascending))) .collect(); - let result = invoke(list(&rows), ascending, true)?; + let result = invoke(list(&rows), ascending, false)?; assert_eq!(bits(&result), expected, "ascending={ascending}"); - let result = invoke(list(&rows).slice(37, 120), ascending, true)?; + let result = invoke(list(&rows).slice(37, 120), ascending, false)?; assert_eq!(bits(&result), &expected[37..157], "ascending={ascending}"); } } @@ -441,7 +401,7 @@ mod tests { /// 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 non_null_ascending_follows_java() -> Result<()> { + fn negative_zero_first_follows_java() -> Result<()> { let long_bits = |v: f64| { if v.is_nan() { f64::NAN.to_bits() as i64 @@ -464,15 +424,14 @@ mod tests { }) }) .collect(); - let result = invoke(list(&rows), true, false)?; + let result = invoke(list(&rows), true, true)?; assert_eq!(bits(&result), expected); Ok(()) } - /// Spark's generated code sorts an ascending array that cannot hold a null with - /// `java.util.Arrays.sort`, which puts -0.0 before 0.0. Descending still ties them. + /// With `-0.0` first, the zeros no longer tie and keep their order. #[test] - fn non_null_ascending_puts_negative_zero_first() -> Result<()> { + fn negative_zero_first() -> Result<()> { let rows = vec![Some(vec![ Some(0.0), Some(1.0), @@ -480,25 +439,26 @@ mod tests { Some(NEGATIVE_NAN), ])]; let cases = [ - (true, false, vec![-0.0, 0.0, 1.0, NEGATIVE_NAN]), - (true, true, vec![0.0, -0.0, 1.0, NEGATIVE_NAN]), - (false, false, vec![NEGATIVE_NAN, 1.0, 0.0, -0.0]), + (true, vec![-0.0, 0.0, 1.0, NEGATIVE_NAN]), + (false, vec![0.0, -0.0, 1.0, NEGATIVE_NAN]), ]; - for (ascending, contains_null, expected) in cases { - let result = invoke(list(&rows), ascending, contains_null)?; + for (negative_zero_first, expected) in cases { + let result = invoke(list(&rows), true, negative_zero_first)?; let expected: Vec> = expected.iter().map(|v| Some(v.to_bits())).collect(); assert_eq!( bits(&result), vec![Some(expected)], - "ascending={ascending} contains_null={contains_null}" + "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 non_list_is_an_error() { - assert!(invoke(Arc::new(Float64Array::from(vec![1.0])), true, true).is_err()); + 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] @@ -512,7 +472,7 @@ mod tests { None, ]), ])); - let result = invoke(array, true, true)?; + let result = invoke(array, true, false)?; let row = result.as_list::().value(0); let row: &Float32Array = row.as_primitive(); let actual: Vec> = row.iter().map(|v| v.map(f32::to_bits)).collect(); @@ -548,17 +508,13 @@ mod tests { (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]), ] { - // Not primitive, so Spark sorts in its SQL ordering whatever containsNull says. let result = invoke(Arc::clone(&array), ascending, false)?; - let row = result.as_list::().value(0); - let x = row - .as_struct() - .column(0) - .as_primitive::() - .clone(); - let actual: Vec = x.values().iter().map(|v| v.to_bits()).collect(); let expected: Vec = expected.iter().map(|v| v.to_bits()).collect(); - assert_eq!(actual, expected, "ascending={ascending}"); + assert_eq!( + first_row_field_bits(&result), + expected, + "ascending={ascending}" + ); } Ok(()) } diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs index 7d79c3dd315..b4ab1ae1308 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(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(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(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/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 6a07a5a36f0..6a5ce015618 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -45,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) } } @@ -182,13 +190,15 @@ object CometSortArray extends CometExpressionSerde[SortArray] with CodegenDispat classOf[FloatType], classOf[DoubleType])) { // DataFusion's array_sort orders floats by IEEE 754 total order. spark_sort_array follows - // Spark, whose generated code also sorts differently when the elements cannot be null, so it - // takes containsNull, which the native array type does not carry. + // 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(arrayType.containsNull), inputs, binding)) + exprToProtoInternal(Literal(negativeZeroFirst), inputs, binding)) } else { val direction = if (ascending) "ASC" else "DESC" val nullOrdering = if (ascending) "NULLS FIRST" else "NULLS LAST" 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 index b7443ab0603..d944d837f4c 100644 --- 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 @@ -22,8 +22,6 @@ -- `-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. --- ConfigMatrix: spark.comet.exec.strictFloatingPoint=false,true - statement CREATE TABLE ar_float(id INT, d DOUBLE, f FLOAT) USING parquet 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 index 39439d25214..95a6f6a8d73 100644 --- 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 @@ -24,10 +24,6 @@ -- `-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. --- Strict floating-point mode no longer needs the codegen dispatcher for these. --- ConfigMatrix: spark.comet.exec.strictFloatingPoint=false,true --- Config: spark.comet.expression.SortArray.allowIncompatible=false - statement CREATE TABLE sa_float(id INT, d DOUBLE, f FLOAT) USING parquet From ce83b32ad3702ed6577bbe1d3b30755a64a69448 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 2 Oct 2026 05:19:45 -0600 Subject: [PATCH 4/4] test: drop the sweep suite's known gap for array_remove and sort_array array_remove and sort_array now compare floats as Spark does, so the suite requires those cases to match. --- .../scala/org/apache/comet/CometFloatSemanticsSuite.scala | 4 ---- 1 file changed, 4 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala b/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala index 1e6aa60fdf1..7901f1004d7 100644 --- a/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala @@ -378,10 +378,6 @@ object CometFloatSemanticsSuite { c => in("expression", "greatest", "least")(c) || in("aggregate", "max", "min", "global max and min", "window max and min")(c)), - KnownGap( - issue(6385), - "array_remove and sort_array compare floats by IEEE 754 total order.", - in("expression", "array_remove", "sort_array", "sort_array descending")), KnownGap( issue(5701), "array_distinct and array_union fold -0.0 into 0.0 and keep NaNs with different bits " +