feat(spark): add equal_null

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 }

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<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 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,
+        )))
+    }
+}

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()]
 }

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;