chore: refactor array fn signatures & add more slt tests (#17672)

Author: Jefffrey

datafusion/functions-nested/src/array_has.rs

@@ -497,7 +497,7 @@ impl Default for ArrayHasAll {
 impl ArrayHasAll {
     pub fn new() -> Self {
         Self {
-            signature: Signature::any(2, Volatility::Immutable),
+            signature: Signature::arrays(2, None, Volatility::Immutable),
             aliases: vec![String::from("list_has_all")],
         }
     }
@@ -571,7 +571,7 @@ impl Default for ArrayHasAny {
 impl ArrayHasAny {
     pub fn new() -> Self {
         Self {
-            signature: Signature::any(2, Volatility::Immutable),
+            signature: Signature::arrays(2, None, Volatility::Immutable),
             aliases: vec![String::from("list_has_any"), String::from("arrays_overlap")],
         }
     }

datafusion/functions-nested/src/cardinality.rs

@@ -58,7 +58,6 @@ impl Cardinality {
                 ],
                 Volatility::Immutable,
             ),
-            aliases: vec![],
         }
     }
 }
@@ -83,7 +82,6 @@ impl Cardinality {
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct Cardinality {
     signature: Signature,
-    aliases: Vec<String>,
 }
 
 impl Default for Cardinality {
@@ -114,10 +112,6 @@ impl ScalarUDFImpl for Cardinality {
         make_scalar_function(cardinality_inner)(&args.args)
     }
 
-    fn aliases(&self) -> &[String] {
-        &self.aliases
-    }
-
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }

datafusion/functions-nested/src/flatten.rs

@@ -25,11 +25,9 @@ use arrow::datatypes::{
     DataType::{FixedSizeList, LargeList, List, Null},
 };
 use datafusion_common::cast::{as_large_list_array, as_list_array};
-use datafusion_common::utils::ListCoercion;
 use datafusion_common::{exec_err, utils::take_function_args, Result};
 use datafusion_expr::{
-    ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
-    ScalarUDFImpl, Signature, TypeSignature, Volatility,
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
 use std::any::Any;
@@ -75,15 +73,7 @@ impl Default for Flatten {
 impl Flatten {
     pub fn new() -> Self {
         Self {
-            signature: Signature {
-                type_signature: TypeSignature::ArraySignature(
-                    ArrayFunctionSignature::Array {
-                        arguments: vec![ArrayFunctionArgument::Array],
-                        array_coercion: Some(ListCoercion::FixedSizedListToList),
-                    },
-                ),
-                volatility: Volatility::Immutable,
-            },
+            signature: Signature::array(Volatility::Immutable),
             aliases: vec![],
         }
     }
@@ -104,7 +94,7 @@ impl ScalarUDFImpl for Flatten {
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         let data_type = match &arg_types[0] {
-            List(field) | FixedSizeList(field, _) => match field.data_type() {
+            List(field) => match field.data_type() {
                 List(field) | FixedSizeList(field, _) => List(Arc::clone(field)),
                 _ => arg_types[0].clone(),
             },

datafusion/functions-nested/src/length.rs

@@ -29,9 +29,10 @@ use arrow::datatypes::{
 use datafusion_common::cast::{
     as_fixed_size_list_array, as_generic_list_array, as_int64_array,
 };
-use datafusion_common::{exec_err, plan_err, Result};
+use datafusion_common::{exec_err, Result};
 use datafusion_expr::{
-    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
+    ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
+    ScalarUDFImpl, Signature, TypeSignature, Volatility,
 };
 use datafusion_functions::downcast_arg;
 use datafusion_macros::user_doc;
@@ -79,7 +80,22 @@ impl Default for ArrayLength {
 impl ArrayLength {
     pub fn new() -> Self {
         Self {
-            signature: Signature::variadic_any(Volatility::Immutable),
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
+                        arguments: vec![ArrayFunctionArgument::Array],
+                        array_coercion: None,
+                    }),
+                    TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
+                        arguments: vec![
+                            ArrayFunctionArgument::Array,
+                            ArrayFunctionArgument::Index,
+                        ],
+                        array_coercion: None,
+                    }),
+                ],
+                Volatility::Immutable,
+            ),
             aliases: vec![String::from("list_length")],
         }
     }
@@ -97,13 +113,8 @@ impl ScalarUDFImpl for ArrayLength {
         &self.signature
     }
 
-    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        Ok(match arg_types[0] {
-            List(_) | LargeList(_) | FixedSizeList(_, _) => UInt64,
-            _ => {
-                return plan_err!("The array_length function can only accept List/LargeList/FixedSizeList.");
-            }
-        })
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(UInt64)
     }
 
     fn invoke_with_args(

datafusion/functions-nested/src/remove.rs

@@ -26,9 +26,11 @@ use arrow::array::{
 use arrow::buffer::OffsetBuffer;
 use arrow::datatypes::{DataType, Field};
 use datafusion_common::cast::as_int64_array;
+use datafusion_common::utils::ListCoercion;
 use datafusion_common::{exec_err, utils::take_function_args, Result};
 use datafusion_expr::{
-    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
+    ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
+    ScalarUDFImpl, Signature, TypeSignature, Volatility,
 };
 use datafusion_macros::user_doc;
 use std::any::Any;
@@ -156,7 +158,17 @@ pub(super) struct ArrayRemoveN {
 impl ArrayRemoveN {
     pub fn new() -> Self {
         Self {
-            signature: Signature::any(3, Volatility::Immutable),
+            signature: Signature::new(
+                TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
+                    arguments: vec![
+                        ArrayFunctionArgument::Array,
+                        ArrayFunctionArgument::Element,
+                        ArrayFunctionArgument::Index,
+                    ],
+                    array_coercion: Some(ListCoercion::FixedSizedListToList),
+                }),
+                Volatility::Immutable,
+            ),
             aliases: vec!["list_remove_n".to_string()],
         }
     }

datafusion/functions-nested/src/resize.rs

@@ -26,7 +26,7 @@ use arrow::buffer::OffsetBuffer;
 use arrow::datatypes::DataType;
 use arrow::datatypes::{ArrowNativeType, Field};
 use arrow::datatypes::{
-    DataType::{FixedSizeList, LargeList, List},
+    DataType::{LargeList, List},
     FieldRef,
 };
 use datafusion_common::cast::{as_int64_array, as_large_list_array, as_list_array};
@@ -125,7 +125,7 @@ impl ScalarUDFImpl for ArrayResize {
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         match &arg_types[0] {
-            List(field) | FixedSizeList(field, _) => Ok(List(Arc::clone(field))),
+            List(field) => Ok(List(Arc::clone(field))),
             LargeList(field) => Ok(LargeList(Arc::clone(field))),
             DataType::Null => {
                 Ok(List(Arc::new(Field::new_list_field(DataType::Int64, true))))

datafusion/functions-nested/src/reverse.rs

@@ -76,7 +76,7 @@ impl Default for ArrayReverse {
 impl ArrayReverse {
     pub fn new() -> Self {
         Self {
-            signature: Signature::any(1, Volatility::Immutable),
+            signature: Signature::array(Volatility::Immutable),
             aliases: vec!["list_reverse".to_string()],
         }
     }