support Signature::Equal

yyy1000 · Apr 28, 2024 · a531ca5 · a531ca5
1 parent f8c623f
commit a531ca5
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 10 deletions.
diff --git a/datafusion/expr/src/signature.rs b/datafusion/expr/src/signature.rs
@@ -108,6 +108,8 @@ pub enum TypeSignature {
     /// 1. A function of one argument of f64 is `Uniform(1, vec![DataType::Float64])`
     /// 2. A function of one argument of f64 or f32 is `Uniform(1, vec![DataType::Float32, DataType::Float64])`
     Uniform(usize, Vec<DataType>),
+    /// Exact number of arguments of for the first argument of a list of valid types.
+    Equal(usize),
     /// Exact number of arguments of an exact type
     Exact(Vec<DataType>),
     /// Fixed number of arguments of arbitrary types
@@ -184,6 +186,12 @@ impl TypeSignature {
             TypeSignature::Exact(types) => {
                 vec![Self::join_types(types, ", ")]
             }
+            TypeSignature::Equal(arg_count) => {
+                vec![std::iter::repeat("Type")
+                    .take(*arg_count)
+                    .collect::<Vec<&str>>()
+                    .join(", ")]
+            }
             TypeSignature::Any(arg_count) => {
                 vec![std::iter::repeat("Any")
                     .take(*arg_count)
@@ -280,6 +288,13 @@ impl Signature {
             volatility,
         }
     }
+    ///
+    pub fn equal(arg_count: usize, volatility: Volatility) -> Self {
+        Self {
+            type_signature: TypeSignature::Equal(arg_count),
+            volatility,
+        }
+    }
     /// Exactly matches the types in `exact_types`, in order.
     pub fn exact(exact_types: Vec<DataType>, volatility: Volatility) -> Self {
         Signature {

diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
@@ -184,6 +184,9 @@ fn get_valid_types(
             .iter()
             .map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())
             .collect(),
+        TypeSignature::Equal(number) => {
+            vec![vec![current_types[0].clone(); *number]]
+        }
         TypeSignature::VariadicEqual => {
             let new_type = current_types.iter().skip(1).try_fold(
                 current_types.first().unwrap().clone(),

diff --git a/datafusion/functions/src/core/nullif.rs b/datafusion/functions/src/core/nullif.rs
@@ -59,11 +59,7 @@ impl Default for NullIfFunc {
 impl NullIfFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::uniform(
-                2,
-                SUPPORTED_NULLIF_TYPES.to_vec(),
-                Volatility::Immutable,
-            ),
+            signature: Signature::equal(2, Volatility::Immutable),
         }
     }
 }

diff --git a/datafusion/functions/src/core/nvl.rs b/datafusion/functions/src/core/nvl.rs
@@ -56,11 +56,7 @@ impl Default for NVLFunc {
 impl NVLFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::uniform(
-                2,
-                SUPPORTED_NVL_TYPES.to_vec(),
-                Volatility::Immutable,
-            ),
+            signature: Signature::equal(2, Volatility::Immutable),
             aliases: vec![String::from("ifnull")],
         }
     }