Fix array_element function signature

datafusion/expr/src/built_in_function.rs

@@ -954,7 +954,7 @@ impl BuiltinScalarFunction {
                 Signature::array(true, self.volatility())
             }
             BuiltinScalarFunction::ArrayElement => {
-                Signature::array_and_index(false, self.volatility())
+                Signature::array_and_index(true, self.volatility())
             }
             BuiltinScalarFunction::ArrayExcept => Signature::any(2, self.volatility()),
             BuiltinScalarFunction::Flatten => Signature::array(true, self.volatility()),

datafusion/expr/src/signature.rs

@@ -122,29 +122,32 @@ pub enum TypeSignature {
     /// is `OneOf(vec![Any(0), VariadicAny])`.
     OneOf(Vec<TypeSignature>),
     /// Specifies Signatures for array functions
-    /// Boolean value specifies whether null type coercion is allowed
-    ArraySignature(ArrayFunctionSignature, bool),
+    ArraySignature(ArrayFunctionSignature),
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum ArrayFunctionSignature {
     /// Specialized Signature for ArrayAppend and similar functions
+    /// If `allow_null` is true, the function also accepts a single argument of type Null.
     /// The first argument should be List/LargeList/FixedSizedList, 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.
     /// List dimension of the List/LargeList is equivalent to the number of List.
     /// List dimension of the non-list is 0.
-    ArrayAndElement,
+    ArrayAndElement(bool),
     /// Specialized Signature for ArrayPrepend and similar functions
+    /// If `allow_null` is true, the function also accepts a single argument of type Null.
     /// 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,
+    ElementAndArray(bool),
     /// Specialized Signature for ArrayEqual and similar functions
+    /// If `allow_null` is true, the function also accepts a single argument of type Null.
     /// The first argument should be List/LargeList/FixedSizedList, and the second argument should be Int64.
-    ArrayAndIndex,
+    ArrayAndIndex(bool),
     /// Specialized Signature for ArrayEmpty and similar functions
     /// The function takes a single argument that must be a List/LargeList/FixedSizeList
     /// or something that can be coerced to one of those types.
-    Array,
+    /// If `allow_null` is true, the function also accepts a single argument of type Null.
+    Array(bool),
 }
 
 impl ArrayFunctionSignature {
@@ -155,7 +158,6 @@ impl ArrayFunctionSignature {
     pub fn get_type_signature(
         &self,
         current_types: &[DataType],
-        allow_null_coercion: bool,
     ) -> Result<Vec<Vec<DataType>>> {
         fn array_append_or_prepend_valid_types(
             current_types: &[DataType],
@@ -208,20 +210,28 @@ impl ArrayFunctionSignature {
                 _ => Ok(vec![vec![]]),
             }
         }
-        fn array_and_index(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
+        fn array_and_index(
+            current_types: &[DataType],
+            allow_null_coercion: bool,
+        ) -> Result<Vec<Vec<DataType>>> {
             if current_types.len() != 2 {
                 return Ok(vec![vec![]]);
             }
 
             let array_type = &current_types[0];
 
+            if array_type.eq(&DataType::Null) && !allow_null_coercion {
+                return Ok(vec![vec![]]);
+            }
+
             match array_type {
                 DataType::List(_)
                 | DataType::LargeList(_)
                 | DataType::FixedSizeList(_, _) => {
                     let array_type = coerced_fixed_size_list_to_list(array_type);
                     Ok(vec![vec![array_type, DataType::Int64]])
                 }
+                DataType::Null => Ok(vec![vec![array_type.clone(), DataType::Int64]]),
                 _ => Ok(vec![vec![]]),
             }
         }
@@ -253,40 +263,36 @@ impl ArrayFunctionSignature {
             }
         }
         match self {
-            ArrayFunctionSignature::ArrayAndElement => {
-                array_append_or_prepend_valid_types(
-                    current_types,
-                    true,
-                    allow_null_coercion,
-                )
+            ArrayFunctionSignature::ArrayAndElement(allow_null) => {
+                array_append_or_prepend_valid_types(current_types, true, *allow_null)
             }
-            ArrayFunctionSignature::ElementAndArray => {
-                array_append_or_prepend_valid_types(
-                    current_types,
-                    false,
-                    allow_null_coercion,
-                )
+            ArrayFunctionSignature::ElementAndArray(allow_null) => {
+                array_append_or_prepend_valid_types(current_types, false, *allow_null)
+            }
+            ArrayFunctionSignature::ArrayAndIndex(allow_null) => {
+                array_and_index(current_types, *allow_null)
+            }
+            ArrayFunctionSignature::Array(allow_null) => {
+                array(current_types, *allow_null)
             }
-            ArrayFunctionSignature::ArrayAndIndex => array_and_index(current_types),
-            ArrayFunctionSignature::Array => array(current_types, allow_null_coercion),
         }
     }
 }
 
 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::ArrayAndElement(allow_null) => {
+                write!(f, "ArrayAndElement({})", *allow_null)
             }
-            ArrayFunctionSignature::ElementAndArray => {
-                write!(f, "element, array")
+            ArrayFunctionSignature::ElementAndArray(allow_null) => {
+                write!(f, "ElementAndArray({})", *allow_null)
             }
-            ArrayFunctionSignature::ArrayAndIndex => {
-                write!(f, "array, index")
+            ArrayFunctionSignature::ArrayAndIndex(allow_null) => {
+                write!(f, "ArrayAndIndex({})", *allow_null)
             }
-            ArrayFunctionSignature::Array => {
-                write!(f, "array")
+            ArrayFunctionSignature::Array(allow_null) => {
+                write!(f, "Array({})", *allow_null)
             }
         }
     }
@@ -320,7 +326,7 @@ impl TypeSignature {
             TypeSignature::OneOf(sigs) => {
                 sigs.iter().flat_map(|s| s.to_string_repr()).collect()
             }
-            TypeSignature::ArraySignature(array_signature, _) => {
+            TypeSignature::ArraySignature(array_signature) => {
                 vec![array_signature.to_string()]
             }
         }
@@ -425,42 +431,38 @@ impl Signature {
         }
     }
     /// Specialized Signature for ArrayAppend and similar functions
-    pub fn array_and_element(allow_null_coercion: bool, volatility: Volatility) -> Self {
+    pub fn array_and_element(allow_null: bool, volatility: Volatility) -> Self {
         Signature {
             type_signature: TypeSignature::ArraySignature(
-                ArrayFunctionSignature::ArrayAndElement,
-                allow_null_coercion,
+                ArrayFunctionSignature::ArrayAndElement(allow_null),
             ),
             volatility,
         }
     }
     /// Specialized Signature for ArrayPrepend and similar functions
-    pub fn element_and_array(allow_null_coercion: bool, volatility: Volatility) -> Self {
+    pub fn element_and_array(allow_null: bool, volatility: Volatility) -> Self {
         Signature {
             type_signature: TypeSignature::ArraySignature(
-                ArrayFunctionSignature::ElementAndArray,
-                allow_null_coercion,
+                ArrayFunctionSignature::ElementAndArray(allow_null),
             ),
             volatility,
         }
     }
     /// Specialized Signature for ArrayElement and similar functions
-    pub fn array_and_index(allow_null_coercion: bool, volatility: Volatility) -> Self {
+    pub fn array_and_index(allow_null: bool, volatility: Volatility) -> Self {
         Signature {
             type_signature: TypeSignature::ArraySignature(
-                ArrayFunctionSignature::ArrayAndIndex,
-                allow_null_coercion,
+                ArrayFunctionSignature::ArrayAndIndex(allow_null),
             ),
             volatility,
         }
     }
     /// Specialized Signature for ArrayEmpty and similar functions
-    pub fn array(allow_null_coercion: bool, volatility: Volatility) -> Self {
+    pub fn array(allow_null: bool, volatility: Volatility) -> Self {
         Signature {
-            type_signature: TypeSignature::ArraySignature(
-                ArrayFunctionSignature::Array,
-                allow_null_coercion,
-            ),
+            type_signature: TypeSignature::ArraySignature(ArrayFunctionSignature::Array(
+                allow_null,
+            )),
             volatility,
         }
     }

datafusion/expr/src/type_coercion/functions.rs

@@ -110,8 +110,8 @@ fn get_valid_types(
         }
 
         TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
-        TypeSignature::ArraySignature(ref function_signature, allow_null_coercion) => {
-            function_signature.get_type_signature(current_types, *allow_null_coercion)?
+        TypeSignature::ArraySignature(ref function_signature) => {
+            function_signature.get_type_signature(current_types)?
         }
 
         TypeSignature::Any(number) => {

datafusion/physical-expr/src/array_expressions.rs

@@ -433,6 +433,7 @@ pub fn array_element(args: &[ArrayRef]) -> Result<ArrayRef> {
             let indexes = as_int64_array(&args[1])?;
             general_array_element::<i64>(array, indexes)
         }
+        DataType::Null => Ok(args[0].clone()),
         _ => exec_err!(
             "array_element does not support type: {:?}",
             args[0].data_type()

datafusion/sqllogictest/test_files/array.slt

@@ -1084,7 +1084,7 @@ from arrays_values_without_nulls;
 ## array_element (aliases: array_extract, list_extract, list_element)
 
 # array_element error
-query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_element\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_element\(array, index\)
+query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_element\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_element\(ArrayAndIndex\(true\)\)
 select array_element(1, 2);
 
 # array_element with null
@@ -4267,7 +4267,7 @@ NULL 10
 ## array_dims (aliases: `list_dims`)
 
 # array dims error
-query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_dims\(Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_dims\(array\)
+query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_dims\(Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_dims\(Array\(false\)\)
 select array_dims(1);
 
 # array_dims scalar function