refactor type signature

apache · Jan 26, 2024 · 2d16386 · 2d16386
1 parent 8c0a92e
commit 2d16386
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 19 deletions.
diff --git 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
@@ -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<TypeSignature>),
+    /// 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>, T)".to_string()]
-            }
-            TypeSignature::ElementAndArray => {
-                vec!["ElementAndArray(T, List<T>)".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
@@ -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!(