From d0f27c642f900a13b0be50aabb2c051abc3689cc Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Mon, 11 Nov 2024 15:15:53 +0800 Subject: [PATCH 1/4] 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), } } } From 825e71412cd081d6eb0e3d8a977059d70118315f Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Tue, 12 Nov 2024 10:25:25 +0800 Subject: [PATCH 2/4] rename to nullary Signed-off-by: jayzhan211 --- datafusion/expr-common/src/signature.rs | 14 +++++++------- datafusion/expr/src/type_coercion/functions.rs | 4 ++-- datafusion/functions-aggregate/src/count.rs | 2 +- datafusion/functions-nested/src/make_array.rs | 2 +- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index ddd361adcb29..bbc5765120b5 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -135,7 +135,7 @@ pub enum TypeSignature { /// Dictionary with string value type is also handled. String(usize), /// Zero argument - ZeroArg, + NullAry, } #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] @@ -192,7 +192,7 @@ impl std::fmt::Display for ArrayFunctionSignature { impl TypeSignature { pub fn to_string_repr(&self) -> Vec { match self { - TypeSignature::ZeroArg => { + TypeSignature::NullAry => { vec!["ZeroArg()".to_string()] } TypeSignature::Variadic(types) => { @@ -248,7 +248,7 @@ impl TypeSignature { pub fn supports_zero_argument(&self) -> bool { match &self { TypeSignature::Exact(vec) => vec.is_empty(), - TypeSignature::ZeroArg => true, + TypeSignature::NullAry => true, TypeSignature::OneOf(types) => types .iter() .any(|type_sig| type_sig.supports_zero_argument()), @@ -291,7 +291,7 @@ impl TypeSignature { .collect(), // TODO: Implement for other types TypeSignature::Any(_) - | TypeSignature::ZeroArg + | TypeSignature::NullAry | TypeSignature::VariadicAny | TypeSignature::ArraySignature(_) | TypeSignature::UserDefined => vec![], @@ -414,7 +414,7 @@ impl Signature { pub fn zero_arg(volatility: Volatility) -> Self { Signature { - type_signature: TypeSignature::ZeroArg, + type_signature: TypeSignature::NullAry, volatility, } } @@ -491,10 +491,10 @@ mod tests { TypeSignature::Exact(vec![]), TypeSignature::OneOf(vec![ TypeSignature::Exact(vec![DataType::Int8]), - TypeSignature::ZeroArg, + TypeSignature::NullAry, TypeSignature::Uniform(1, vec![DataType::Int8]), ]), - TypeSignature::ZeroArg, + TypeSignature::NullAry, ]; for case in positive_cases { diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 997089cdbb41..6836713d8016 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -181,7 +181,7 @@ fn is_well_supported_signature(type_signature: &TypeSignature) -> bool { | TypeSignature::String(_) | TypeSignature::Coercible(_) | TypeSignature::Any(_) - | TypeSignature::ZeroArg + | TypeSignature::NullAry ) } @@ -618,7 +618,7 @@ fn get_valid_types( } } }, - TypeSignature::ZeroArg => { + TypeSignature::NullAry => { if !current_types.is_empty() { return plan_err!( "The function expected zero argument but received {}", diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 2001746177de..52181372698f 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -102,7 +102,7 @@ impl Count { pub fn new() -> Self { Self { signature: Signature::one_of( - vec![TypeSignature::VariadicAny, TypeSignature::ZeroArg], + vec![TypeSignature::VariadicAny, TypeSignature::NullAry], Volatility::Immutable, ), } diff --git a/datafusion/functions-nested/src/make_array.rs b/datafusion/functions-nested/src/make_array.rs index b6d23b6fd988..de67b0ae3874 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::ZeroArg, TypeSignature::UserDefined], + vec![TypeSignature::NullAry, TypeSignature::UserDefined], Volatility::Immutable, ), aliases: vec![String::from("make_list")], From 26abab9913709b1053090825a883a90d3571ee9c Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Tue, 12 Nov 2024 13:51:08 +0800 Subject: [PATCH 3/4] rename Signed-off-by: jayzhan211 --- datafusion/expr-common/src/signature.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 +- datafusion/functions/src/datetime/current_date.rs | 2 +- datafusion/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 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index bbc5765120b5..1746333d4a06 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -412,7 +412,7 @@ impl Signature { } } - pub fn zero_arg(volatility: Volatility) -> Self { + pub fn nullary(volatility: Volatility) -> Self { Signature { type_signature: TypeSignature::NullAry, volatility, diff --git a/datafusion/functions-window/src/cume_dist.rs b/datafusion/functions-window/src/cume_dist.rs index d575e7cd28e8..500d96b56323 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::zero_arg(Volatility::Immutable), + signature: Signature::nullary(Volatility::Immutable), } } } diff --git a/datafusion/functions-window/src/rank.rs b/datafusion/functions-window/src/rank.rs index 1bbbb6ebb9dd..06945e693eea 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::zero_arg(Volatility::Immutable), + signature: Signature::nullary(Volatility::Immutable), rank_type, } } diff --git a/datafusion/functions-window/src/row_number.rs b/datafusion/functions-window/src/row_number.rs index bf056b4c9609..68f6fde23280 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::zero_arg(Volatility::Immutable), + signature: Signature::nullary(Volatility::Immutable), } } } diff --git a/datafusion/functions/src/datetime/current_date.rs b/datafusion/functions/src/datetime/current_date.rs index bfb10303d66e..3b819c470d1e 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::zero_arg(Volatility::Stable), + signature: Signature::nullary(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 630dd3692d87..ca591f922305 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::zero_arg(Volatility::Stable), + signature: Signature::nullary(Volatility::Stable), } } } diff --git a/datafusion/functions/src/datetime/now.rs b/datafusion/functions/src/datetime/now.rs index 87f62d17f67c..cadc4fce04f1 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::zero_arg(Volatility::Stable), + signature: Signature::nullary(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 21b0eace056a..70cc76f03c58 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::zero_arg(Volatility::Immutable), + signature: Signature::nullary(Volatility::Immutable), } } } diff --git a/datafusion/functions/src/math/random.rs b/datafusion/functions/src/math/random.rs index ed5987e4ca03..0026037c95bd 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::zero_arg(Volatility::Volatile), + signature: Signature::nullary(Volatility::Volatile), } } } From d02297ce642d6a3d278a25c22f3404542462c332 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Tue, 12 Nov 2024 15:11:55 +0800 Subject: [PATCH 4/4] tostring Signed-off-by: jayzhan211 --- datafusion/expr-common/src/signature.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 1746333d4a06..0fffd84b7047 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -193,7 +193,7 @@ impl TypeSignature { pub fn to_string_repr(&self) -> Vec { match self { TypeSignature::NullAry => { - vec!["ZeroArg()".to_string()] + vec!["NullAry()".to_string()] } TypeSignature::Variadic(types) => { vec![format!("{}, ..", Self::join_types(types, "/"))]