simplify signature for nullif

Signed-off-by: Jay Zhan <[email protected]>

Author: jayzhan211

datafusion/expr-common/src/signature.rs

@@ -133,6 +133,15 @@ pub enum TypeSignature {
     /// Null is considerd as `Utf8` by default
     /// Dictionary with string value type is also handled.
     String(usize),
+    /// Fixed number of arguments of boolean types.
+    Boolean(usize),
+}
+
+impl TypeSignature {
+    #[inline]
+    pub fn is_one_of(&self) -> bool {
+        matches!(self, TypeSignature::OneOf(_))
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
@@ -198,6 +207,9 @@ impl TypeSignature {
                     .collect::<Vec<String>>()
                     .join(", ")]
             }
+            TypeSignature::Boolean(num) => {
+                vec![format!("Boolean({num})")]
+            }
             TypeSignature::String(num) => {
                 vec![format!("String({num})")]
             }
@@ -267,6 +279,7 @@ impl TypeSignature {
             | TypeSignature::UserDefined
             | TypeSignature::ArraySignature(_)
             | TypeSignature::Numeric(_)
+            | TypeSignature::Boolean(_)
             | TypeSignature::String(_) => vec![],
         }
     }
@@ -323,6 +336,14 @@ impl Signature {
         }
     }
 
+    /// A specified number of boolean arguments
+    pub fn boolean(arg_count: usize, volatility: Volatility) -> Self {
+        Self {
+            type_signature: TypeSignature::Boolean(arg_count),
+            volatility,
+        }
+    }
+
     /// An arbitrary number of arguments of any type.
     pub fn variadic_any(volatility: Volatility) -> Self {
         Self {

datafusion/expr/src/type_coercion/functions.rs

@@ -180,6 +180,7 @@ fn is_well_supported_signature(type_signature: &TypeSignature) -> bool {
             | TypeSignature::Numeric(_)
             | TypeSignature::String(_)
             | TypeSignature::Coercible(_)
+            | TypeSignature::Boolean(_)
             | TypeSignature::Any(_)
     )
 }
@@ -193,13 +194,18 @@ fn try_coerce_types(
 
     // Well-supported signature that returns exact valid types.
     if !valid_types.is_empty() && is_well_supported_signature(type_signature) {
-        // exact valid types
-        assert_eq!(valid_types.len(), 1);
+        // There may be many valid types if valid signature is OneOf
+        // Otherwise, there should be only one valid type
+        if !type_signature.is_one_of() {
+            assert_eq!(valid_types.len(), 1);
+        }
+
         let valid_types = valid_types.swap_remove(0);
         if let Some(t) = maybe_data_types_without_coercion(&valid_types, current_types) {
             return Ok(t);
         }
     } else {
+        // TODO: Deprecate this branch after all signatures are well-supported (aka coercion are happend already)
         // Try and coerce the argument types to match the signature, returning the
         // coerced types from the first matching signature.
         for valid_types in valid_types {
@@ -476,6 +482,33 @@ fn get_valid_types(
 
             vec![vec![base_type_or_default_type(&coerced_type); *number]]
         }
+        TypeSignature::Boolean(number) => {
+            function_length_check(current_types.len(), *number)?;
+
+            // Find common boolean type amongs given types
+            let mut valid_type = current_types.first().unwrap().to_owned();
+            for t in current_types.iter().skip(1) {
+                let logical_data_type: NativeType = t.into();
+                if logical_data_type == NativeType::Null {
+                    continue;
+                }
+
+                if logical_data_type != NativeType::Boolean {
+                    return plan_err!(
+                        "The signature expected NativeType::Boolean but received {logical_data_type}"
+                    );
+                }
+
+                valid_type = t.to_owned();
+            }
+
+            let logical_data_type: NativeType = valid_type.clone().into();
+            if logical_data_type == NativeType::Null {
+                valid_type = DataType::Boolean;
+            }
+
+            vec![vec![valid_type; *number]]
+        }
         TypeSignature::Numeric(number) => {
             function_length_check(current_types.len(), *number)?;
 

datafusion/functions/src/core/nullif.rs

@@ -17,7 +17,7 @@
 
 use arrow::datatypes::DataType;
 use datafusion_common::{exec_err, Result};
-use datafusion_expr::{ColumnarValue, Documentation};
+use datafusion_expr::{ColumnarValue, Documentation, TypeSignature};
 
 use arrow::compute::kernels::cmp::eq;
 use arrow::compute::kernels::nullif::nullif;
@@ -32,25 +32,6 @@ pub struct NullIfFunc {
     signature: Signature,
 }
 
-/// Currently supported types by the nullif function.
-/// The order of these types correspond to the order on which coercion applies
-/// This should thus be from least informative to most informative
-static SUPPORTED_NULLIF_TYPES: &[DataType] = &[
-    DataType::Boolean,
-    DataType::UInt8,
-    DataType::UInt16,
-    DataType::UInt32,
-    DataType::UInt64,
-    DataType::Int8,
-    DataType::Int16,
-    DataType::Int32,
-    DataType::Int64,
-    DataType::Float32,
-    DataType::Float64,
-    DataType::Utf8,
-    DataType::LargeUtf8,
-];
-
 impl Default for NullIfFunc {
     fn default() -> Self {
         Self::new()
@@ -60,9 +41,13 @@ impl Default for NullIfFunc {
 impl NullIfFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::uniform(
-                2,
-                SUPPORTED_NULLIF_TYPES.to_vec(),
+            signature: Signature::one_of(
+                // Hack: String is at the beginning so the return type is String if both args are Nulls
+                vec![
+                    TypeSignature::String(2),
+                    TypeSignature::Numeric(2),
+                    TypeSignature::Boolean(2),
+                ],
                 Volatility::Immutable,
             ),
         }
@@ -82,14 +67,7 @@ impl ScalarUDFImpl for NullIfFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        // NULLIF has two args and they might get coerced, get a preview of this
-        let coerced_types = datafusion_expr::type_coercion::functions::data_types(
-            arg_types,
-            &self.signature,
-        );
-        coerced_types
-            .map(|typs| typs[0].clone())
-            .map_err(|e| e.context("Failed to coerce arguments for NULLIF"))
+        Ok(arg_types[0].to_owned())
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {

datafusion/sqllogictest/test_files/nullif.slt

@@ -97,7 +97,31 @@ SELECT NULLIF(1, 3);
 ----
 1
 
-query I
+query T
 SELECT NULLIF(NULL, NULL);
 ----
 NULL
+
+query R
+select nullif(1, 1.2);
+----
+1
+
+query R
+select nullif(1.0, 2);
+----
+1
+
+query error DataFusion error: Error during planning: Internal error: Failed to match any signature, errors: Error during planning: The signature expected NativeType::String but received NativeType::Int64
+select nullif(2, 'a');
+
+
+query T
+select nullif('2', '3');
+----
+2
+
+# TODO: support numeric string
+# This query success in Postgres and DuckDB
+query error DataFusion error: Error during planning: Internal error: Failed to match any signature, errors: Error during planning: The signature expected NativeType::String but received NativeType::Int64
+select nullif(2, '1');