Skip to content

Commit 8d6899e

Browse files
authored
Support TypeSignature::Nullary (#13354)
* support zero arg Signed-off-by: jayzhan211 <[email protected]> * rename to nullary Signed-off-by: jayzhan211 <[email protected]> * rename Signed-off-by: jayzhan211 <[email protected]> * tostring Signed-off-by: jayzhan211 <[email protected]> --------- Signed-off-by: jayzhan211 <[email protected]>
1 parent 1bbe13f commit 8d6899e

File tree

12 files changed

+58
-21
lines changed

12 files changed

+58
-21
lines changed

datafusion/expr-common/src/signature.rs

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,7 @@ pub enum TypeSignature {
113113
/// arguments like `vec![DataType::Int32]` or `vec![DataType::Float32]`
114114
/// since i32 and f32 can be casted to f64
115115
Coercible(Vec<LogicalTypeRef>),
116-
/// Fixed number of arguments of arbitrary types
117-
/// If a function takes 0 argument, its `TypeSignature` should be `Any(0)`
116+
/// Fixed number of arguments of arbitrary types, number should be larger than 0
118117
Any(usize),
119118
/// Matches exactly one of a list of [`TypeSignature`]s. Coercion is attempted to match
120119
/// the signatures in order, and stops after the first success, if any.
@@ -135,6 +134,8 @@ pub enum TypeSignature {
135134
/// Null is considerd as `Utf8` by default
136135
/// Dictionary with string value type is also handled.
137136
String(usize),
137+
/// Zero argument
138+
NullAry,
138139
}
139140

140141
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
@@ -191,6 +192,9 @@ impl std::fmt::Display for ArrayFunctionSignature {
191192
impl TypeSignature {
192193
pub fn to_string_repr(&self) -> Vec<String> {
193194
match self {
195+
TypeSignature::NullAry => {
196+
vec!["NullAry()".to_string()]
197+
}
194198
TypeSignature::Variadic(types) => {
195199
vec![format!("{}, ..", Self::join_types(types, "/"))]
196200
}
@@ -244,7 +248,7 @@ impl TypeSignature {
244248
pub fn supports_zero_argument(&self) -> bool {
245249
match &self {
246250
TypeSignature::Exact(vec) => vec.is_empty(),
247-
TypeSignature::Uniform(0, _) | TypeSignature::Any(0) => true,
251+
TypeSignature::NullAry => true,
248252
TypeSignature::OneOf(types) => types
249253
.iter()
250254
.any(|type_sig| type_sig.supports_zero_argument()),
@@ -287,6 +291,7 @@ impl TypeSignature {
287291
.collect(),
288292
// TODO: Implement for other types
289293
TypeSignature::Any(_)
294+
| TypeSignature::NullAry
290295
| TypeSignature::VariadicAny
291296
| TypeSignature::ArraySignature(_)
292297
| TypeSignature::UserDefined => vec![],
@@ -407,6 +412,13 @@ impl Signature {
407412
}
408413
}
409414

415+
pub fn nullary(volatility: Volatility) -> Self {
416+
Signature {
417+
type_signature: TypeSignature::NullAry,
418+
volatility,
419+
}
420+
}
421+
410422
/// A specified number of arguments of any type
411423
pub fn any(arg_count: usize, volatility: Volatility) -> Self {
412424
Signature {
@@ -477,13 +489,12 @@ mod tests {
477489
// Testing `TypeSignature`s which supports 0 arg
478490
let positive_cases = vec![
479491
TypeSignature::Exact(vec![]),
480-
TypeSignature::Uniform(0, vec![DataType::Float64]),
481-
TypeSignature::Any(0),
482492
TypeSignature::OneOf(vec![
483493
TypeSignature::Exact(vec![DataType::Int8]),
484-
TypeSignature::Any(0),
494+
TypeSignature::NullAry,
485495
TypeSignature::Uniform(1, vec![DataType::Int8]),
486496
]),
497+
TypeSignature::NullAry,
487498
];
488499

489500
for case in positive_cases {

datafusion/expr/src/type_coercion/functions.rs

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,7 @@ fn is_well_supported_signature(type_signature: &TypeSignature) -> bool {
181181
| TypeSignature::String(_)
182182
| TypeSignature::Coercible(_)
183183
| TypeSignature::Any(_)
184+
| TypeSignature::NullAry
184185
)
185186
}
186187

@@ -554,16 +555,27 @@ fn get_valid_types(
554555

555556
vec![new_types]
556557
}
557-
TypeSignature::Uniform(number, valid_types) => valid_types
558-
.iter()
559-
.map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())
560-
.collect(),
558+
TypeSignature::Uniform(number, valid_types) => {
559+
if *number == 0 {
560+
return plan_err!("The function expected at least one argument");
561+
}
562+
563+
valid_types
564+
.iter()
565+
.map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())
566+
.collect()
567+
}
561568
TypeSignature::UserDefined => {
562569
return internal_err!(
563570
"User-defined signature should be handled by function-specific coerce_types."
564571
)
565572
}
566573
TypeSignature::VariadicAny => {
574+
if current_types.is_empty() {
575+
return plan_err!(
576+
"The function expected at least one argument but received 0"
577+
);
578+
}
567579
vec![current_types.to_vec()]
568580
}
569581
TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
@@ -606,7 +618,22 @@ fn get_valid_types(
606618
}
607619
}
608620
},
621+
TypeSignature::NullAry => {
622+
if !current_types.is_empty() {
623+
return plan_err!(
624+
"The function expected zero argument but received {}",
625+
current_types.len()
626+
);
627+
}
628+
vec![vec![]]
629+
}
609630
TypeSignature::Any(number) => {
631+
if current_types.is_empty() {
632+
return plan_err!(
633+
"The function expected at least one argument but received 0"
634+
);
635+
}
636+
610637
if current_types.len() != *number {
611638
return plan_err!(
612639
"The function expected {} arguments but received {}",

datafusion/functions-aggregate/src/count.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,7 @@ impl Count {
102102
pub fn new() -> Self {
103103
Self {
104104
signature: Signature::one_of(
105-
// TypeSignature::Any(0) is required to handle `Count()` with no args
106-
vec![TypeSignature::VariadicAny, TypeSignature::Any(0)],
105+
vec![TypeSignature::VariadicAny, TypeSignature::NullAry],
107106
Volatility::Immutable,
108107
),
109108
}

datafusion/functions-nested/src/make_array.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ impl MakeArray {
6363
pub fn new() -> Self {
6464
Self {
6565
signature: Signature::one_of(
66-
vec![TypeSignature::UserDefined, TypeSignature::Any(0)],
66+
vec![TypeSignature::NullAry, TypeSignature::UserDefined],
6767
Volatility::Immutable,
6868
),
6969
aliases: vec![String::from("make_list")],

datafusion/functions-window/src/cume_dist.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ pub struct CumeDist {
4949
impl CumeDist {
5050
pub fn new() -> Self {
5151
Self {
52-
signature: Signature::any(0, Volatility::Immutable),
52+
signature: Signature::nullary(Volatility::Immutable),
5353
}
5454
}
5555
}

datafusion/functions-window/src/rank.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ impl Rank {
7474
pub fn new(name: String, rank_type: RankType) -> Self {
7575
Self {
7676
name,
77-
signature: Signature::any(0, Volatility::Immutable),
77+
signature: Signature::nullary(Volatility::Immutable),
7878
rank_type,
7979
}
8080
}

datafusion/functions-window/src/row_number.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ impl RowNumber {
5151
/// Create a new `row_number` function
5252
pub fn new() -> Self {
5353
Self {
54-
signature: Signature::any(0, Volatility::Immutable),
54+
signature: Signature::nullary(Volatility::Immutable),
5555
}
5656
}
5757
}

datafusion/functions/src/datetime/current_date.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ impl Default for CurrentDateFunc {
4444
impl CurrentDateFunc {
4545
pub fn new() -> Self {
4646
Self {
47-
signature: Signature::uniform(0, vec![], Volatility::Stable),
47+
signature: Signature::nullary(Volatility::Stable),
4848
aliases: vec![String::from("today")],
4949
}
5050
}

datafusion/functions/src/datetime/current_time.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ impl Default for CurrentTimeFunc {
4242
impl CurrentTimeFunc {
4343
pub fn new() -> Self {
4444
Self {
45-
signature: Signature::uniform(0, vec![], Volatility::Stable),
45+
signature: Signature::nullary(Volatility::Stable),
4646
}
4747
}
4848
}

datafusion/functions/src/datetime/now.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ impl Default for NowFunc {
4343
impl NowFunc {
4444
pub fn new() -> Self {
4545
Self {
46-
signature: Signature::uniform(0, vec![], Volatility::Stable),
46+
signature: Signature::nullary(Volatility::Stable),
4747
aliases: vec!["current_timestamp".to_string()],
4848
}
4949
}

0 commit comments

Comments
 (0)