support zero arg

Signed-off-by: jayzhan211 <[email protected]>
jayzhan211 · Nov 11, 2024 · d0f27c6 · d0f27c6
1 parent 7ebd993
commit d0f27c6
Show file tree

Hide file tree

Showing 12 changed files with 58 additions and 21 deletions.
diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs
@@ -113,8 +113,7 @@ pub enum TypeSignature {
     /// arguments like `vec![DataType::Int32]` or `vec![DataType::Float32]`
     /// since i32 and f32 can be casted to f64
     Coercible(Vec<LogicalTypeRef>),
-    /// Fixed number of arguments of arbitrary types
-    /// If a function takes 0 argument, its `TypeSignature` should be `Any(0)`
+    /// Fixed number of arguments of arbitrary types, number should be larger than 0
     Any(usize),
     /// Matches exactly one of a list of [`TypeSignature`]s. Coercion is attempted to match
     /// the signatures in order, and stops after the first success, if any.
@@ -135,6 +134,8 @@ pub enum TypeSignature {
     /// Null is considerd as `Utf8` by default
     /// Dictionary with string value type is also handled.
     String(usize),
+    /// Zero argument
+    ZeroArg,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
@@ -191,6 +192,9 @@ impl std::fmt::Display for ArrayFunctionSignature {
 impl TypeSignature {
     pub fn to_string_repr(&self) -> Vec<String> {
         match self {
+            TypeSignature::ZeroArg => {
+                vec!["ZeroArg()".to_string()]
+            }
             TypeSignature::Variadic(types) => {
                 vec![format!("{}, ..", Self::join_types(types, "/"))]
             }
@@ -244,7 +248,7 @@ impl TypeSignature {
     pub fn supports_zero_argument(&self) -> bool {
         match &self {
             TypeSignature::Exact(vec) => vec.is_empty(),
-            TypeSignature::Uniform(0, _) | TypeSignature::Any(0) => true,
+            TypeSignature::ZeroArg => true,
             TypeSignature::OneOf(types) => types
                 .iter()
                 .any(|type_sig| type_sig.supports_zero_argument()),
@@ -287,6 +291,7 @@ impl TypeSignature {
                 .collect(),
             // TODO: Implement for other types
             TypeSignature::Any(_)
+            | TypeSignature::ZeroArg
             | TypeSignature::VariadicAny
             | TypeSignature::ArraySignature(_)
             | TypeSignature::UserDefined => vec![],
@@ -407,6 +412,13 @@ impl Signature {
         }
     }
 
+    pub fn zero_arg(volatility: Volatility) -> Self {
+        Signature {
+            type_signature: TypeSignature::ZeroArg,
+            volatility,
+        }
+    }
+
     /// A specified number of arguments of any type
     pub fn any(arg_count: usize, volatility: Volatility) -> Self {
         Signature {
@@ -477,13 +489,12 @@ mod tests {
         // Testing `TypeSignature`s which supports 0 arg
         let positive_cases = vec![
             TypeSignature::Exact(vec![]),
-            TypeSignature::Uniform(0, vec![DataType::Float64]),
-            TypeSignature::Any(0),
             TypeSignature::OneOf(vec![
                 TypeSignature::Exact(vec![DataType::Int8]),
-                TypeSignature::Any(0),
+                TypeSignature::ZeroArg,
                 TypeSignature::Uniform(1, vec![DataType::Int8]),
             ]),
+            TypeSignature::ZeroArg,
         ];
 
         for case in positive_cases {

diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
@@ -181,6 +181,7 @@ fn is_well_supported_signature(type_signature: &TypeSignature) -> bool {
             | TypeSignature::String(_)
             | TypeSignature::Coercible(_)
             | TypeSignature::Any(_)
+            | TypeSignature::ZeroArg
     )
 }
 
@@ -554,16 +555,27 @@ fn get_valid_types(
 
             vec![new_types]
         }
-        TypeSignature::Uniform(number, valid_types) => valid_types
-            .iter()
-            .map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())
-            .collect(),
+        TypeSignature::Uniform(number, valid_types) => {
+            if *number == 0 {
+                return plan_err!("The function expected at least one argument");
+            }
+
+            valid_types
+                .iter()
+                .map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())
+                .collect()
+        }
         TypeSignature::UserDefined => {
             return internal_err!(
             "User-defined signature should be handled by function-specific coerce_types."
         )
         }
         TypeSignature::VariadicAny => {
+            if current_types.is_empty() {
+                return plan_err!(
+                    "The function expected at least one argument but received 0"
+                );
+            }
             vec![current_types.to_vec()]
         }
         TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
@@ -606,7 +618,22 @@ fn get_valid_types(
                 }
             }
         },
+        TypeSignature::ZeroArg => {
+            if !current_types.is_empty() {
+                return plan_err!(
+                    "The function expected zero argument but received {}",
+                    current_types.len()
+                );
+            }
+            vec![vec![]]
+        }
         TypeSignature::Any(number) => {
+            if current_types.is_empty() {
+                return plan_err!(
+                    "The function expected at least one argument but received 0"
+                );
+            }
+
             if current_types.len() != *number {
                 return plan_err!(
                     "The function expected {} arguments but received {}",

diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs
@@ -102,8 +102,7 @@ impl Count {
     pub fn new() -> Self {
         Self {
             signature: Signature::one_of(
-                // TypeSignature::Any(0) is required to handle `Count()` with no args
-                vec![TypeSignature::VariadicAny, TypeSignature::Any(0)],
+                vec![TypeSignature::VariadicAny, TypeSignature::ZeroArg],
                 Volatility::Immutable,
             ),
         }

diff --git a/datafusion/functions-nested/src/make_array.rs b/datafusion/functions-nested/src/make_array.rs
@@ -63,7 +63,7 @@ impl MakeArray {
     pub fn new() -> Self {
         Self {
             signature: Signature::one_of(
-                vec![TypeSignature::UserDefined, TypeSignature::Any(0)],
+                vec![TypeSignature::ZeroArg, TypeSignature::UserDefined],
                 Volatility::Immutable,
             ),
             aliases: vec![String::from("make_list")],

diff --git a/datafusion/functions-window/src/cume_dist.rs b/datafusion/functions-window/src/cume_dist.rs
@@ -49,7 +49,7 @@ pub struct CumeDist {
 impl CumeDist {
     pub fn new() -> Self {
         Self {
-            signature: Signature::any(0, Volatility::Immutable),
+            signature: Signature::zero_arg(Volatility::Immutable),
         }
     }
 }

diff --git a/datafusion/functions-window/src/rank.rs b/datafusion/functions-window/src/rank.rs
@@ -74,7 +74,7 @@ impl Rank {
     pub fn new(name: String, rank_type: RankType) -> Self {
         Self {
             name,
-            signature: Signature::any(0, Volatility::Immutable),
+            signature: Signature::zero_arg(Volatility::Immutable),
             rank_type,
         }
     }

diff --git a/datafusion/functions-window/src/row_number.rs b/datafusion/functions-window/src/row_number.rs
@@ -51,7 +51,7 @@ impl RowNumber {
     /// Create a new `row_number` function
     pub fn new() -> Self {
         Self {
-            signature: Signature::any(0, Volatility::Immutable),
+            signature: Signature::zero_arg(Volatility::Immutable),
         }
     }
 }

diff --git a/datafusion/functions/src/datetime/current_date.rs b/datafusion/functions/src/datetime/current_date.rs
@@ -44,7 +44,7 @@ impl Default for CurrentDateFunc {
 impl CurrentDateFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::uniform(0, vec![], Volatility::Stable),
+            signature: Signature::zero_arg(Volatility::Stable),
             aliases: vec![String::from("today")],
         }
     }

diff --git a/datafusion/functions/src/datetime/current_time.rs b/datafusion/functions/src/datetime/current_time.rs
@@ -42,7 +42,7 @@ impl Default for CurrentTimeFunc {
 impl CurrentTimeFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::uniform(0, vec![], Volatility::Stable),
+            signature: Signature::zero_arg(Volatility::Stable),
         }
     }
 }

diff --git a/datafusion/functions/src/datetime/now.rs b/datafusion/functions/src/datetime/now.rs
@@ -43,7 +43,7 @@ impl Default for NowFunc {
 impl NowFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::uniform(0, vec![], Volatility::Stable),
+            signature: Signature::zero_arg(Volatility::Stable),
             aliases: vec!["current_timestamp".to_string()],
         }
     }

diff --git a/datafusion/functions/src/math/pi.rs b/datafusion/functions/src/math/pi.rs
@@ -41,7 +41,7 @@ impl Default for PiFunc {
 impl PiFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::exact(vec![], Volatility::Immutable),
+            signature: Signature::zero_arg(Volatility::Immutable),
         }
     }
 }

diff --git a/datafusion/functions/src/math/random.rs b/datafusion/functions/src/math/random.rs
@@ -42,7 +42,7 @@ impl Default for RandomFunc {
 impl RandomFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::exact(vec![], Volatility::Volatile),
+            signature: Signature::zero_arg(Volatility::Volatile),
         }
     }
 }