Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support TypeSignature::Nullary #13354

Merged
merged 4 commits into from
Nov 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions datafusion/expr-common/src/signature.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
NullAry,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the A here intended to be upper-case? I'd have expected Nullary

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nullary looks good to me too

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
Expand Down Expand Up @@ -191,6 +192,9 @@ impl std::fmt::Display for ArrayFunctionSignature {
impl TypeSignature {
pub fn to_string_repr(&self) -> Vec<String> {
match self {
TypeSignature::NullAry => {
vec!["NullAry()".to_string()]
}
TypeSignature::Variadic(types) => {
vec![format!("{}, ..", Self::join_types(types, "/"))]
}
Expand Down Expand Up @@ -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::NullAry => true,
TypeSignature::OneOf(types) => types
.iter()
.any(|type_sig| type_sig.supports_zero_argument()),
Expand Down Expand Up @@ -287,6 +291,7 @@ impl TypeSignature {
.collect(),
// TODO: Implement for other types
TypeSignature::Any(_)
| TypeSignature::NullAry
| TypeSignature::VariadicAny
| TypeSignature::ArraySignature(_)
| TypeSignature::UserDefined => vec![],
Expand Down Expand Up @@ -407,6 +412,13 @@ impl Signature {
}
}

pub fn nullary(volatility: Volatility) -> Self {
Signature {
type_signature: TypeSignature::NullAry,
volatility,
}
}

/// A specified number of arguments of any type
pub fn any(arg_count: usize, volatility: Volatility) -> Self {
Signature {
Expand Down Expand Up @@ -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::NullAry,
TypeSignature::Uniform(1, vec![DataType::Int8]),
]),
TypeSignature::NullAry,
];

for case in positive_cases {
Expand Down
35 changes: 31 additions & 4 deletions datafusion/expr/src/type_coercion/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@ fn is_well_supported_signature(type_signature: &TypeSignature) -> bool {
| TypeSignature::String(_)
| TypeSignature::Coercible(_)
| TypeSignature::Any(_)
| TypeSignature::NullAry
)
}

Expand Down Expand Up @@ -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()],
Expand Down Expand Up @@ -606,7 +618,22 @@ fn get_valid_types(
}
}
},
TypeSignature::NullAry => {
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 {}",
Expand Down
3 changes: 1 addition & 2 deletions datafusion/functions-aggregate/src/count.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::NullAry],
Volatility::Immutable,
),
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-nested/src/make_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ impl MakeArray {
pub fn new() -> Self {
Self {
signature: Signature::one_of(
vec![TypeSignature::UserDefined, TypeSignature::Any(0)],
vec![TypeSignature::NullAry, TypeSignature::UserDefined],
Volatility::Immutable,
),
aliases: vec![String::from("make_list")],
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-window/src/cume_dist.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ pub struct CumeDist {
impl CumeDist {
pub fn new() -> Self {
Self {
signature: Signature::any(0, Volatility::Immutable),
signature: Signature::nullary(Volatility::Immutable),
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-window/src/rank.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::nullary(Volatility::Immutable),
rank_type,
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-window/src/row_number.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::nullary(Volatility::Immutable),
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions/src/datetime/current_date.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ impl Default for CurrentDateFunc {
impl CurrentDateFunc {
pub fn new() -> Self {
Self {
signature: Signature::uniform(0, vec![], Volatility::Stable),
signature: Signature::nullary(Volatility::Stable),
aliases: vec![String::from("today")],
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions/src/datetime/current_time.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ impl Default for CurrentTimeFunc {
impl CurrentTimeFunc {
pub fn new() -> Self {
Self {
signature: Signature::uniform(0, vec![], Volatility::Stable),
signature: Signature::nullary(Volatility::Stable),
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions/src/datetime/now.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ impl Default for NowFunc {
impl NowFunc {
pub fn new() -> Self {
Self {
signature: Signature::uniform(0, vec![], Volatility::Stable),
signature: Signature::nullary(Volatility::Stable),
aliases: vec!["current_timestamp".to_string()],
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions/src/math/pi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ impl Default for PiFunc {
impl PiFunc {
pub fn new() -> Self {
Self {
signature: Signature::exact(vec![], Volatility::Immutable),
signature: Signature::nullary(Volatility::Immutable),
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions/src/math/random.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ impl Default for RandomFunc {
impl RandomFunc {
pub fn new() -> Self {
Self {
signature: Signature::exact(vec![], Volatility::Volatile),
signature: Signature::nullary(Volatility::Volatile),
}
}
}
Expand Down