Update array_slice signature

datafusion/expr-common/src/signature.rs

@@ -19,6 +19,7 @@
 //! and return types of functions in DataFusion.
 
 use std::fmt::Display;
+use std::num::NonZeroUsize;
 
 use crate::type_coercion::aggregates::NUMERICS;
 use arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
@@ -236,9 +237,9 @@ pub enum ArrayFunctionSignature {
     /// 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,
-    /// Specialized Signature for Array functions of the form (List/LargeList, Index)
-    /// The first argument should be List/LargeList/FixedSizedList, and the second argument should be Int64.
-    ArrayAndIndex,
+    /// Specialized Signature for Array functions of the form (List/LargeList, Index+)
+    /// The first argument should be List/LargeList/FixedSizedList, and the next n arguments should be Int64.
+    ArrayAndIndexes(NonZeroUsize),
     /// Specialized Signature for Array functions of the form (List/LargeList, Element, Optional Index)
     ArrayAndElementAndOptionalIndex,
     /// Specialized Signature for ArrayEmpty and similar functions
@@ -265,8 +266,12 @@ impl Display for ArrayFunctionSignature {
             ArrayFunctionSignature::ElementAndArray => {
                 write!(f, "element, array")
             }
-            ArrayFunctionSignature::ArrayAndIndex => {
-                write!(f, "array, index")
+            ArrayFunctionSignature::ArrayAndIndexes(count) => {
+                write!(f, "array")?;
+                for _ in 0..count.get() {
+                    write!(f, ", index")?;
+                }
+                Ok(())
             }
             ArrayFunctionSignature::Array => {
                 write!(f, "array")
@@ -627,9 +632,13 @@ impl Signature {
     }
     /// Specialized Signature for ArrayElement and similar functions
     pub fn array_and_index(volatility: Volatility) -> Self {
+        Self::array_and_indexes(volatility, NonZeroUsize::new(1).expect("1 is non-zero"))
+    }
+    /// Specialized Signature for ArraySlice and similar functions
+    pub fn array_and_indexes(volatility: Volatility, count: NonZeroUsize) -> Self {
         Signature {
             type_signature: TypeSignature::ArraySignature(
-                ArrayFunctionSignature::ArrayAndIndex,
+                ArrayFunctionSignature::ArrayAndIndexes(count),
             ),
             volatility,
             null_handling: NullHandling::PassThrough,

datafusion/expr/src/type_coercion/functions.rs

@@ -671,13 +671,20 @@ fn get_valid_types(
             ArrayFunctionSignature::ElementAndArray => {
                 array_append_or_prepend_valid_types(current_types, false)?
             }
-            ArrayFunctionSignature::ArrayAndIndex => {
-                if current_types.len() != 2 {
+            ArrayFunctionSignature::ArrayAndIndexes(count) => {
+                if current_types.len() != count.get() + 1 {
                     return Ok(vec![vec![]]);
                 }
                 array(&current_types[0]).map_or_else(
                     || vec![vec![]],
-                    |array_type| vec![vec![array_type, DataType::Int64]],
+                    |array_type| {
+                        let mut inner = Vec::with_capacity(count.get() + 1);
+                        inner.push(array_type);
+                        for _ in 0..count.get() {
+                            inner.push(DataType::Int64);
+                        }
+                        vec![inner]
+                    },
                 )
             }
             ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {

datafusion/functions-nested/src/extract.rs

@@ -35,12 +35,13 @@ use datafusion_common::cast::as_list_array;
 use datafusion_common::{
     exec_err, internal_datafusion_err, plan_err, DataFusionError, Result,
 };
-use datafusion_expr::Expr;
+use datafusion_expr::{ArrayFunctionSignature, Expr, TypeSignature};
 use datafusion_expr::{
     ColumnarValue, Documentation, NullHandling, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
 use std::any::Any;
+use std::num::NonZeroUsize;
 use std::sync::Arc;
 
 use crate::utils::make_scalar_function;
@@ -330,9 +331,27 @@ pub(super) struct ArraySlice {
 impl ArraySlice {
     pub fn new() -> Self {
         Self {
-            // TODO: This signature should use the actual accepted types, not variadic_any.
-            signature: Signature::variadic_any(Volatility::Immutable)
-                .with_null_handling(NullHandling::Propagate),
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::ArraySignature(
+                        ArrayFunctionSignature::ArrayAndIndexes(
+                            NonZeroUsize::new(1).expect("1 is non-zero"),
+                        ),
+                    ),
+                    TypeSignature::ArraySignature(
+                        ArrayFunctionSignature::ArrayAndIndexes(
+                            NonZeroUsize::new(2).expect("2 is non-zero"),
+                        ),
+                    ),
+                    TypeSignature::ArraySignature(
+                        ArrayFunctionSignature::ArrayAndIndexes(
+                            NonZeroUsize::new(3).expect("3 is non-zero"),
+                        ),
+                    ),
+                ],
+                Volatility::Immutable,
+            )
+            .with_null_handling(NullHandling::Propagate),
             aliases: vec![String::from("list_slice")],
         }
     }

datafusion/functions-nested/src/flatten.rs

@@ -26,7 +26,10 @@ use datafusion_common::cast::{
     as_generic_list_array, as_large_list_array, as_list_array,
 };
 use datafusion_common::{exec_err, Result};
-use datafusion_expr::{ArrayFunctionSignature, ColumnarValue, Documentation, NullHandling, ScalarUDFImpl, Signature, TypeSignature, Volatility};
+use datafusion_expr::{
+    ArrayFunctionSignature, ColumnarValue, Documentation, NullHandling, ScalarUDFImpl,
+    Signature, TypeSignature, Volatility,
+};
 use datafusion_macros::user_doc;
 use std::any::Any;
 use std::sync::Arc;

datafusion/physical-expr/src/scalar_function.rs

@@ -188,7 +188,9 @@ impl PhysicalExpr for ScalarFunctionExpr {
             .collect::<Result<Vec<_>>>()?;
 
         if self.fun.signature().null_handling == NullHandling::Propagate
-            && args.iter().any(|arg| arg.data_type().is_null())
+            && args.iter().any(
+                |arg| matches!(arg, ColumnarValue::Scalar(scalar) if scalar.is_null()),
+            )
         {
             let null_value = ScalarValue::try_from(&self.return_type)?;
             return Ok(ColumnarValue::Scalar(null_value));

datafusion/sqllogictest/test_files/array.slt

@@ -2055,6 +2055,9 @@ select array_slice(a, -1, 2, 1), array_slice(a, -1, 2),
 query error DataFusion error: Error during planning: array_slice does not support zero arguments
 select array_slice();
 
+query error Failed to coerce arguments
+select array_slice(3.5, NULL, NULL);
+
 ## array_any_value (aliases: list_any_value)
 
 # Testing with empty arguments should result in an error