This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24827-a948ff63807a10dd61c3d20fd78f84ebb3e71461 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit d42cd854df0f78d95e95946831077db3c7808ced Author: Advit Arora <[email protected]> AuthorDate: Tue Sep 22 08:48:42 2026 +0000 feat(spark): add equal_null (#24827) ## Which issue does this PR close? No issue filed. Part of #15914, which tracks the Spark function library. ## Rationale for this change `equal_null` is missing, and the `misc` module it belongs in had no functions in it at all. In Spark it is an alias: `EqualNull(l, r)` is rewritten to `EqualNullSafe`, the `<=>` operator, so it is null-safe equality. Both NULL is true, and the result is never NULL: ```sql SELECT equal_null(NULL, NULL); -- true SELECT equal_null(NULL::int, 1::int); -- false ``` DataFusion already has that operator as `IS NOT DISTINCT FROM`, so this mirrors Spark's own structure instead of writing a second comparison kernel. ## What changes are included in this PR? `misc/equal_null.rs` and its registration. `simplify()` rewrites to the operator, which is Spark's `replacement` field, and `invoke_with_args` calls `apply_cmp` with the same operator so the function also works with the logical optimizer disabled, as the crate README requires for Comet. Two of the stub's commented-out queries were malformed, the porting script wrote one cast per distinct `typeof()` key, so a call with two identical literals lost an argument. The same artifact affects 10 more pairs in 6 other spark files, left alone here. Two divergences from Spark are left alone because they belong to the operator, not to this function. Spark treats `-NaN` and `NaN` as equal and DataFusion does not, which is true of `IS NOT DISTINCT FROM` generally. Spark also rejects maps at analysis since `MapType` is not orderable, while `comparison_coercion` here accepts them. `return_field_from_args` is overridden so the field is not nullable. `normalize_float_zero` now recurses into nested children. `compare_op_for_nested` and `GroupValuesColumn` both normalize through it, so the scalar kernels and the two group-by paths agree on nested float keys. ## Are these changes tested? Yes, `spark/misc/equal_null.slt` goes from a skipped stub to 27 assertions covering the truth table, float ordering, columns, arrays, structs, decimals and the arity errors. Reverting `simplify()` to plain `Eq` fails 9 of them, so the file is not passing on constant folding. Putting `invoke_with_args` back to a stub fails the two queries that run with the optimizer off, and nothing else. ## Are there any user-facing changes? Yes. `equal_null` is new in `datafusion-spark`, and `-0.0` now matches `+0.0` inside nested values, in comparisons and in `GROUP BY`, matching the scalar kernels. --- Cargo.lock | 1 + datafusion/common/src/utils/mod.rs | 42 ++++- datafusion/physical-expr-common/src/datum.rs | 15 +- .../aggregates/group_values/multi_group_by/mod.rs | 7 + datafusion/spark/Cargo.toml | 1 + datafusion/spark/src/function/misc/equal_null.rs | 99 ++++++++++ datafusion/spark/src/function/misc/mod.rs | 17 +- .../sqllogictest/test_files/negative_zero.slt | 25 +++ .../test_files/spark/misc/equal_null.slt | 208 ++++++++++++++++++++- 9 files changed, 400 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bc8f497624..2f2ea87c91 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2677,6 +2677,7 @@ dependencies = [ "datafusion-functions-aggregate", "datafusion-functions-aggregate-common", "datafusion-functions-nested", + "datafusion-physical-expr-common", "datafusion-session", "log", "num-traits", diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index bf87b9a788..84a8a295a6 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -1447,7 +1447,7 @@ fn fsl_values_row_number(list_size: i32, array_len: usize) -> Result<Int32Array> /// OR-reduction) decides whether to fall through to the rewriting path. /// Only arrays that actually contain `-0.0` pay for a new buffer. pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { - use arrow::array::{Float16Array, Float32Array, Float64Array}; + use arrow::array::{Float16Array, Float32Array, Float64Array, make_array}; use arrow::datatypes::{Float16Type, Float32Type, Float64Type}; // -0.0 has only the sign bit set; no other finite or NaN value shares // this bit pattern, so a strict-equality scan reliably gates the rewrite. @@ -1499,10 +1499,50 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { }); Arc::new(normalized) } + dt if has_float_leaf(dt) => { + let data = array.to_data(); + let children = data + .child_data() + .iter() + .map(|child| normalize_float_zero(&make_array(child.clone())).to_data()) + .collect::<Vec<_>>(); + if children + .iter() + .zip(data.child_data()) + .all(|(new, old)| new.ptr_eq(old)) + { + return Arc::clone(array); + } + make_array( + data.into_builder() + .child_data(children) + .build() + .expect("rewriting float leaves preserves the array layout"), + ) + } _ => Arc::clone(array), } } +pub fn has_float_leaf(data_type: &DataType) -> bool { + match data_type { + DataType::Float16 | DataType::Float32 | DataType::Float64 => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) + | DataType::RunEndEncoded(_, f) => has_float_leaf(f.data_type()), + DataType::Struct(fields) => fields.iter().any(|f| has_float_leaf(f.data_type())), + DataType::Union(fields, _) => { + fields.iter().any(|(_, f)| has_float_leaf(f.data_type())) + } + DataType::Dictionary(_, values) => has_float_leaf(values), + _ => false, + } +} + /// Replace `-0.0` with `+0.0` in `Float16`, `Float32`, or `Float64` scalar /// values. Other variants are returned unchanged. See [`normalize_float_zero`] /// for context. diff --git a/datafusion/physical-expr-common/src/datum.rs b/datafusion/physical-expr-common/src/datum.rs index d23fb30db6..416507c618 100644 --- a/datafusion/physical-expr-common/src/datum.rs +++ b/datafusion/physical-expr-common/src/datum.rs @@ -16,14 +16,16 @@ // under the License. use arrow::array::BooleanArray; -use arrow::array::{ArrayRef, Datum, make_comparator}; +use arrow::array::{Array, ArrayRef, Datum, make_array, make_comparator}; use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::compute::kernels::cmp::{ distinct, eq, gt, gt_eq, lt, lt_eq, neq, not_distinct, }; use arrow::compute::{SortOptions, ilike, like, nilike, nlike}; use arrow::error::ArrowError; -use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; +use datafusion_common::utils::{ + has_float_leaf, normalize_float_zero, normalize_float_zero_scalar, +}; use datafusion_common::{Result, ScalarValue}; use datafusion_common::{arrow_datafusion_err, assert_or_internal_err, internal_err}; use datafusion_expr_common::columnar_value::ColumnarValue; @@ -147,6 +149,11 @@ pub fn compare_with_eq( } } +fn normalize_nested_float_zero(array: &dyn Array) -> Option<ArrayRef> { + has_float_leaf(array.data_type()) + .then(|| normalize_float_zero(&make_array(array.to_data()))) +} + /// Compare on nested type List, Struct, and so on pub fn compare_op_for_nested( op: Operator, @@ -155,6 +162,10 @@ pub fn compare_op_for_nested( ) -> Result<BooleanArray> { let (l, is_l_scalar) = lhs.get(); let (r, is_r_scalar) = rhs.get(); + let l_norm = normalize_nested_float_zero(l); + let r_norm = normalize_nested_float_zero(r); + let l = l_norm.as_deref().unwrap_or(l); + let r = r_norm.as_deref().unwrap_or(r); let l_len = l.len(); let r_len = r.len(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 5c33abd41d..f10990a73a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -50,6 +50,7 @@ use arrow::datatypes::{ }; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; +use datafusion_common::utils::{has_float_leaf, normalize_float_zero}; use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::{EmitTo, GroupSelection}; @@ -1131,6 +1132,12 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> { impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> { fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + let normalized: Option<Vec<ArrayRef>> = cols + .iter() + .any(|col| col.data_type().is_nested() && has_float_leaf(col.data_type())) + .then(|| cols.iter().map(normalize_float_zero).collect()); + let cols = normalized.as_deref().unwrap_or(cols); + // `try_new` and the reset points in `emit` / `clear_shrink` keep // `self.group_values` populated with one builder per schema field, // so no lazy initialization is needed here. diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index f40708863d..08309a931d 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -56,6 +56,7 @@ datafusion-functions = { workspace = true } datafusion-functions-aggregate = { workspace = true } datafusion-functions-aggregate-common = { workspace = true } datafusion-functions-nested = { workspace = true } +datafusion-physical-expr-common = { workspace = true } datafusion-session = { workspace = true } log = { workspace = true } num-traits = { workspace = true } diff --git a/datafusion/spark/src/function/misc/equal_null.rs b/datafusion/spark/src/function/misc/equal_null.rs new file mode 100644 index 0000000000..b1604a7162 --- /dev/null +++ b/datafusion/spark/src/function/misc/equal_null.rs @@ -0,0 +1,99 @@ +// 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 arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, plan_err}; +use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; +use datafusion_expr::type_coercion::binary::comparison_coercion; +use datafusion_expr::{ + ColumnarValue, Expr, Operator, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, Volatility, binary_expr, +}; +use datafusion_physical_expr_common::datum::apply_cmp; +use std::sync::Arc; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkEqualNull { + signature: Signature, +} + +impl Default for SparkEqualNull { + fn default() -> Self { + Self::new() + } +} + +impl SparkEqualNull { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkEqualNull { + fn name(&self) -> &str { + "equal_null" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { + let [lhs, rhs] = arg_types else { + return plan_err!( + "Function 'equal_null' expects 2 arguments but received {}", + arg_types.len() + ); + }; + // simplify() emits a comparison, and the type coercion pass has already run by then + let Some(common) = comparison_coercion(lhs, rhs) else { + return plan_err!( + "For function 'equal_null' {lhs} and {rhs} are not comparable" + ); + }; + Ok(vec![common.clone(), common]) + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Boolean) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> { + Ok(Arc::new(Field::new(self.name(), DataType::Boolean, false))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let [lhs, rhs] = take_function_args(self.name(), args.args)?; + apply_cmp(Operator::IsNotDistinctFrom, &lhs, &rhs) + } + + fn simplify( + &self, + args: Vec<Expr>, + _info: &SimplifyContext, + ) -> Result<ExprSimplifyResult> { + let [lhs, rhs] = take_function_args(self.name(), args)?; + Ok(ExprSimplifyResult::Simplified(binary_expr( + lhs, + Operator::IsNotDistinctFrom, + rhs, + ))) + } +} diff --git a/datafusion/spark/src/function/misc/mod.rs b/datafusion/spark/src/function/misc/mod.rs index a87df9a2c8..8739a53f40 100644 --- a/datafusion/spark/src/function/misc/mod.rs +++ b/datafusion/spark/src/function/misc/mod.rs @@ -16,10 +16,23 @@ // under the License. use datafusion_expr::ScalarUDF; +use datafusion_functions::make_udf_function; use std::sync::Arc; -pub mod expr_fn {} +mod equal_null; + +make_udf_function!(equal_null::SparkEqualNull, equal_null); + +pub mod expr_fn { + use datafusion_functions::export_functions; + + export_functions!(( + equal_null, + "Returns true if arg1 equals arg2, or if both are NULL; false otherwise", + arg1 arg2 + )); +} pub fn functions() -> Vec<Arc<ScalarUDF>> { - vec![] + vec![equal_null()] } diff --git a/datafusion/sqllogictest/test_files/negative_zero.slt b/datafusion/sqllogictest/test_files/negative_zero.slt index 8ea1122880..fe4df57e17 100644 --- a/datafusion/sqllogictest/test_files/negative_zero.slt +++ b/datafusion/sqllogictest/test_files/negative_zero.slt @@ -229,3 +229,28 @@ JOIN (SELECT arrow_cast(-0.0, 'Float32') AS b) t2 ON t1.a = t2.b; statement ok reset datafusion.optimizer.prefer_hash_join; + +##### +## Nested values holding +0.0 / -0.0 +##### + +query BBB +SELECT [0.0] = [-0.0] AS eq, + [0.0] IS DISTINCT FROM [-0.0] AS is_distinct, + {a: 0.0} = {a: -0.0} AS struct_eq; +---- +true false true + +statement ok +CREATE TABLE nested_zeros(id INT, a DOUBLE[]) AS VALUES (1, [0.0]), (2, [-0.0]); + +query II +SELECT l.id, r.id FROM nested_zeros l JOIN nested_zeros r ON l.a = r.a ORDER BY l.id, r.id; +---- +1 1 +1 2 +2 1 +2 2 + +statement ok +DROP TABLE nested_zeros; diff --git a/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt b/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt index 71a3af6070..1a12668f60 100644 --- a/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt +++ b/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt @@ -23,25 +23,213 @@ ## Original Query: SELECT equal_null(1, '11'); ## PySpark 3.5.5 Result: {'equal_null(1, 11)': False, 'typeof(equal_null(1, 11))': 'boolean', 'typeof(1)': 'int', 'typeof(11)': 'string'} -#query -#SELECT equal_null(1::int, '11'::string); +query B +SELECT equal_null(1::int, '11'::string); +---- +false ## Original Query: SELECT equal_null(3, 3); ## PySpark 3.5.5 Result: {'equal_null(3, 3)': True, 'typeof(equal_null(3, 3))': 'boolean', 'typeof(3)': 'int'} -#query -#SELECT equal_null(3::int); +query B +SELECT equal_null(3::int, 3::int); +---- +true ## Original Query: SELECT equal_null(NULL, 'abc'); ## PySpark 3.5.5 Result: {'equal_null(NULL, abc)': False, 'typeof(equal_null(NULL, abc))': 'boolean', 'typeof(NULL)': 'void', 'typeof(abc)': 'string'} -#query -#SELECT equal_null(NULL::void, 'abc'::string); +query B +SELECT equal_null(NULL, 'abc'::string); +---- +false ## Original Query: SELECT equal_null(NULL, NULL); ## PySpark 3.5.5 Result: {'equal_null(NULL, NULL)': True, 'typeof(equal_null(NULL, NULL))': 'boolean', 'typeof(NULL)': 'void'} -#query -#SELECT equal_null(NULL::void); +query B +SELECT equal_null(NULL, NULL); +---- +true ## Original Query: SELECT equal_null(true, NULL); ## PySpark 3.5.5 Result: {'equal_null(true, NULL)': False, 'typeof(equal_null(true, NULL))': 'boolean', 'typeof(true)': 'boolean', 'typeof(NULL)': 'void'} -#query -#SELECT equal_null(true::boolean, NULL::void); +query B +SELECT equal_null(true::boolean, NULL); +---- +false + +query B +SELECT equal_null(NULL, true::boolean); +---- +false + +query BB +SELECT equal_null(1::int, 1::int), equal_null(1::int, 2::int); +---- +true false + +query BB +SELECT equal_null(NULL::int, 1::int), equal_null(NULL::int, NULL::int); +---- +false true + +# EqualNullSafe is declared non-nullable in Spark, so the result is never NULL +query B +SELECT equal_null(NULL::int, NULL::int) IS NULL; +---- +false + +query BB +SELECT equal_null('abc'::string, 'abc'::string), equal_null('abc'::string, 'abd'::string); +---- +true false + +# The default UTF8_BINARY collation compares strings by byte +query B +SELECT equal_null('a'::string, 'A'::string); +---- +false + +query BB +SELECT equal_null(true, true), equal_null(true, false); +---- +true false + +query BB +SELECT equal_null(1::int, 1::bigint), equal_null(1::int, 1.0::double); +---- +true true + +# Spark's float ordering makes NaN equal to itself, unlike IEEE-754 +query BB +SELECT equal_null('NaN'::double, 'NaN'::double) AS d, equal_null('NaN'::float, 'NaN'::float) AS f; +---- +true true + +query BBB +SELECT equal_null('NaN'::double, 1.0::double), equal_null('NaN'::double, NULL), equal_null('NaN'::double, 'Infinity'::double); +---- +false false false + +# Spark's float ordering also makes -0.0 equal to 0.0 +query BB +SELECT equal_null(0.0::double, -0.0::double) AS d, equal_null(0.0::float, -0.0::float) AS f; +---- +true true + +query BB +SELECT equal_null('Infinity'::double, 'Infinity'::double), equal_null('Infinity'::double, '-Infinity'::double); +---- +true false + +statement ok +CREATE TABLE equal_null_ints(id INT, a INT, b INT) AS VALUES +(1, 1, 1), +(2, 1, 2), +(3, CAST(NULL AS INT), 1), +(4, 1, CAST(NULL AS INT)), +(5, CAST(NULL AS INT), CAST(NULL AS INT)); + +query B +SELECT equal_null(a, b) FROM equal_null_ints ORDER BY id; +---- +true +false +false +false +true + +statement ok +DROP TABLE equal_null_ints; + +statement ok +CREATE TABLE equal_null_doubles(id INT, a DOUBLE, b DOUBLE) AS VALUES +(1, 'NaN'::double, 'NaN'::double), +(2, 0.0, -0.0), +(3, 1.0, CAST(NULL AS DOUBLE)), +(4, CAST(NULL AS DOUBLE), CAST(NULL AS DOUBLE)); + +query B +SELECT equal_null(a, b) FROM equal_null_doubles ORDER BY id; +---- +true +true +false +true + +statement ok +DROP TABLE equal_null_doubles; + +query BB +SELECT equal_null(array(1, 2), array(1, 2)), equal_null(array(1, 2), array(1, 2, 3)); +---- +true false + +# Two NULLs in the same array slot compare equal, per Spark's array ordering +query BB +SELECT equal_null(array(1, NULL), array(1, NULL)), equal_null(array(1, NULL), array(1, 2)); +---- +true false + +query B +SELECT equal_null(named_struct('a', 1), named_struct('a', 1)); +---- +true + +query BB +SELECT equal_null(array(0.0::double), array(-0.0::double)), equal_null(named_struct('a', 0.0::double), named_struct('a', -0.0::double)); +---- +true true + +query B +SELECT equal_null(1.0::decimal(2,1), 1.00::decimal(3,2)); +---- +true + +statement error Function 'equal_null' expects 2 arguments but received 1 +SELECT equal_null(1::int); + +statement error Function 'equal_null' expects 2 arguments but received 3 +SELECT equal_null(1::int, 2::int, 3::int); + +# Without the simplify() rewrite the function runs its own kernel, which Comet relies on +statement ok +set datafusion.optimizer.max_passes = 0; + +query BBBBB +SELECT equal_null(NULL, NULL), equal_null(0.0::double, -0.0::double), equal_null('NaN'::double, 'NaN'::double), equal_null(array(1, NULL), array(1, NULL)), equal_null(array(0.0::double), array(-0.0::double)); +---- +true true true true true + +statement ok +CREATE TABLE equal_null_physical(id INT, a INT, b INT) AS VALUES +(1, 1, 1), +(2, 1, CAST(NULL AS INT)), +(3, CAST(NULL AS INT), CAST(NULL AS INT)); + +query B +SELECT equal_null(a, b) FROM equal_null_physical ORDER BY id; +---- +true +false +true + +statement ok +DROP TABLE equal_null_physical; + +statement ok +set datafusion.explain.show_schema = true; + +query TT +EXPLAIN SELECT equal_null(NULL::int, NULL::int); +---- +logical_plan +01)Projection: equal_null(CAST(NULL AS Int32), CAST(NULL AS Int32)) +02)--EmptyRelation: rows=1 +physical_plan +01)ProjectionExec: expr=[equal_null(CAST(NULL AS Int32), CAST(NULL AS Int32)) as equal_null(NULL,NULL)], schema=[equal_null(NULL,NULL):Boolean] +02)--PlaceholderRowExec, schema=[] + +statement ok +reset datafusion.explain.show_schema; + +statement ok +reset datafusion.optimizer.max_passes; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
