From 2d16386c4380501a6135f65ca8f0976ef45bfd6a Mon Sep 17 00:00:00 2001 From: Weijun-H Date: Fri, 26 Jan 2024 20:25:27 +0800 Subject: [PATCH] refactor type signature --- datafusion/expr/src/built_in_function.rs | 9 ++-- datafusion/expr/src/signature.rs | 47 ++++++++++++++++--- .../expr/src/type_coercion/functions.rs | 19 +++++--- 3 files changed, 56 insertions(+), 19 deletions(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index 85bf9885ebba..4fef283c56f7 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -945,10 +945,9 @@ impl BuiltinScalarFunction { BuiltinScalarFunction::ArraySort => { Signature::variadic_any(self.volatility()) } - BuiltinScalarFunction::ArrayAppend => Signature { - type_signature: ArrayAndElement, - volatility: self.volatility(), - }, + BuiltinScalarFunction::ArrayAppend => { + Signature::array_and_element(self.volatility()) + } BuiltinScalarFunction::MakeArray => { // 0 or more arguments of arbitrary type Signature::one_of(vec![VariadicEqual, Any(0)], self.volatility()) @@ -961,7 +960,7 @@ impl BuiltinScalarFunction { BuiltinScalarFunction::ArrayDims => Signature::any(1, self.volatility()), BuiltinScalarFunction::ArrayEmpty => Signature::any(1, self.volatility()), BuiltinScalarFunction::ArrayElement => { - Signature::array_and_element(self.volatility()) + Signature::array_and_index(self.volatility()) } BuiltinScalarFunction::ArrayExcept => Signature::any(2, self.volatility()), BuiltinScalarFunction::Flatten => Signature::any(1, self.volatility()), diff --git a/datafusion/expr/src/signature.rs b/datafusion/expr/src/signature.rs index 17a777551323..48f4c996cb5d 100644 --- a/datafusion/expr/src/signature.rs +++ b/datafusion/expr/src/signature.rs @@ -116,6 +116,12 @@ pub enum TypeSignature { /// Function `make_array` takes 0 or more arguments with arbitrary types, its `TypeSignature` /// is `OneOf(vec![Any(0), VariadicAny])`. OneOf(Vec), + /// Specifies Signatures for array functions + ArraySignature(ArrayFunctionSignature), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum ArrayFunctionSignature { /// Specialized Signature for ArrayAppend and similar functions /// The first argument should be List/LargeList, and the second argument should be non-list or list. /// The second argument's list dimension should be one dimension less than the first argument's list dimension. @@ -126,6 +132,23 @@ pub enum TypeSignature { /// The first argument should be non-list or list, and the second argument should be List/LargeList. /// The first argument's list dimension should be one dimension less than the second argument's list dimension. ElementAndArray, + ArrayAndIndex, +} + +impl std::fmt::Display for ArrayFunctionSignature { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ArrayFunctionSignature::ArrayAndElement => { + write!(f, "array, element") + } + ArrayFunctionSignature::ElementAndArray => { + write!(f, "element, array") + } + ArrayFunctionSignature::ArrayAndIndex => { + write!(f, "array, index") + } + } + } } impl TypeSignature { @@ -156,11 +179,8 @@ impl TypeSignature { TypeSignature::OneOf(sigs) => { sigs.iter().flat_map(|s| s.to_string_repr()).collect() } - TypeSignature::ArrayAndElement => { - vec!["ArrayAndElement(List, T)".to_string()] - } - TypeSignature::ElementAndArray => { - vec!["ElementAndArray(T, List)".to_string()] + TypeSignature::ArraySignature(array_signature) => { + vec![array_signature.to_string()] } } } @@ -266,14 +286,27 @@ impl Signature { /// Specialized Signature for ArrayAppend and similar functions pub fn array_and_element(volatility: Volatility) -> Self { Signature { - type_signature: TypeSignature::ArrayAndElement, + type_signature: TypeSignature::ArraySignature( + ArrayFunctionSignature::ArrayAndElement, + ), volatility, } } /// Specialized Signature for ArrayPrepend and similar functions pub fn element_and_array(volatility: Volatility) -> Self { Signature { - type_signature: TypeSignature::ElementAndArray, + type_signature: TypeSignature::ArraySignature( + ArrayFunctionSignature::ElementAndArray, + ), + volatility, + } + } + /// Specialized Signature for ArrayElement and similar functions + pub fn array_and_index(volatility: Volatility) -> Self { + Signature { + type_signature: TypeSignature::ArraySignature( + ArrayFunctionSignature::ArrayAndIndex, + ), volatility, } } diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index ea324c37c30a..e9aa1a813ab1 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::signature::TIMEZONE_WILDCARD; +use crate::signature::{ArrayFunctionSignature, TIMEZONE_WILDCARD}; use crate::{Signature, TypeSignature}; use arrow::{ compute::can_cast_types, @@ -166,12 +166,17 @@ fn get_valid_types( } TypeSignature::Exact(valid_types) => vec![valid_types.clone()], - TypeSignature::ArrayAndElement => { - return array_append_or_prepend_valid_types(current_types, true) - } - TypeSignature::ElementAndArray => { - return array_append_or_prepend_valid_types(current_types, false) - } + TypeSignature::ArraySignature(ref function_signature) => match function_signature + { + ArrayFunctionSignature::ArrayAndElement + | ArrayFunctionSignature::ArrayAndIndex => { + return array_append_or_prepend_valid_types(current_types, true) + } + ArrayFunctionSignature::ElementAndArray => { + return array_append_or_prepend_valid_types(current_types, false) + } + }, + TypeSignature::Any(number) => { if current_types.len() != *number { return plan_err!(