From d0f27c642f900a13b0be50aabb2c051abc3689cc Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Mon, 11 Nov 2024 15:15:53 +0800 Subject: [PATCH] support zero arg Signed-off-by: jayzhan211 --- datafusion/expr-common/src/signature.rs | 23 ++++++++---- .../expr/src/type_coercion/functions.rs | 35 ++++++++++++++++--- datafusion/functions-aggregate/src/count.rs | 3 +- datafusion/functions-nested/src/make_array.rs | 2 +- datafusion/functions-window/src/cume_dist.rs | 2 +- datafusion/functions-window/src/rank.rs | 2 +- datafusion/functions-window/src/row_number.rs | 2 +- .../functions/src/datetime/current_date.rs | 2 +- .../functions/src/datetime/current_time.rs | 2 +- datafusion/functions/src/datetime/now.rs | 2 +- datafusion/functions/src/math/pi.rs | 2 +- datafusion/functions/src/math/random.rs | 2 +- 12 files changed, 58 insertions(+), 21 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 3846fae5de5d..ddd361adcb29 100644 --- 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), - /// 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 { 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 index 5a4d89a0b2ec..997089cdbb41 100644 --- 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 index bade589a908a..2001746177de 100644 --- 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 index 7aa3445f6858..b6d23b6fd988 100644 --- 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 index 9e30c672fee5..d575e7cd28e8 100644 --- 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 index 06c3f49055a5..1bbbb6ebb9dd 100644 --- 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 index 56af14fb84ae..bf056b4c9609 100644 --- 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 index 24046611a71f..bfb10303d66e 100644 --- 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 index 4122b54b07e8..630dd3692d87 100644 --- 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 index c13bbfb18105..87f62d17f67c 100644 --- 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 index 502429d0ca5d..21b0eace056a 100644 --- 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 index cd92798d67dd..ed5987e4ca03 100644 --- 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), } } }