diff --git a/Cargo.lock b/Cargo.lock index d19ce6e86d18e..7fdbcece93f08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2629,6 +2629,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/spark/Cargo.toml b/datafusion/spark/Cargo.toml index f40708863d108..08309a931d96b 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 0000000000000..c6e84cce2d077 --- /dev/null +++ b/datafusion/spark/src/function/misc/equal_null.rs @@ -0,0 +1,94 @@ +// 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; +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, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, binary_expr, +}; +use datafusion_physical_expr_common::datum::apply_cmp; + +#[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> { + 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 { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [lhs, rhs] = take_function_args(self.name(), args.args)?; + apply_cmp(Operator::IsNotDistinctFrom, &lhs, &rhs) + } + + fn simplify( + &self, + args: Vec, + _info: &SimplifyContext, + ) -> Result { + 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 a87df9a2c87a0..8739a53f40adc 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> { - vec![] + vec![equal_null()] } diff --git a/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt b/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt index 88999d997d2db..bc6d6f3f0e7f1 100644 --- a/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt +++ b/datafusion/sqllogictest/test_files/spark/misc/equal_null.slt @@ -23,25 +23,192 @@ ## 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 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 BBBB +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)); +---- +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 +reset datafusion.optimizer.max_passes;