feat: support FixedSizeList Type Coercion

datafusion/expr/src/signature.rs

@@ -236,33 +236,15 @@ impl ArrayFunctionSignature {
                 _ => Ok(vec![vec![]]),
             }
         }
-        fn array_and_index(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
-            if current_types.len() != 2 {
-                return Ok(vec![vec![]]);
-            }
-            let array_type = &current_types[0];
-            let array_type = array(&[array_type.clone()])?;
-            if array_type[0].is_empty() {
-                Ok(vec![vec![]])
-            } else {
-                Ok(vec![vec![array_type[0][0].clone(), DataType::Int64]])
-            }
-        }
-        fn array(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
-            if current_types.len() != 1 {
-                return Ok(vec![vec![]]);
-            }
-
-            let array_type = &current_types[0];
-
+        fn array(array_type: &DataType) -> Option<DataType> {
             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]])
+                    Some(array_type)
                 }
-                _ => Ok(vec![vec![]]),
+                _ => None,
             }
         }
         match self {
@@ -272,11 +254,28 @@ impl ArrayFunctionSignature {
             ArrayFunctionSignature::ElementAndArray => {
                 array_append_or_prepend_valid_types(current_types, false)
             }
-            ArrayFunctionSignature::ArrayAndIndex => array_and_index(current_types),
+            ArrayFunctionSignature::ArrayAndIndex => {
+                if current_types.len() != 2 {
+                    return Ok(vec![vec![]]);
+                }
+                array(&current_types[0]).map_or_else(
+                    || Ok(vec![vec![]]),
+                    |array_type| Ok(vec![vec![array_type, DataType::Int64]]),
+                )
+            }
             ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {
                 array_element_and_optional_index(current_types)
             }
-            ArrayFunctionSignature::Array => array(current_types),
+            ArrayFunctionSignature::Array => {
+                if current_types.len() != 1 {
+                    return Ok(vec![vec![]]);
+                }
+
+                array(&current_types[0]).map_or_else(
+                    || Ok(vec![vec![]]),
+                    |array_type| Ok(vec![vec![array_type]]),
+                )
+            }
         }
     }
 }
@@ -285,19 +284,19 @@ impl std::fmt::Display for ArrayFunctionSignature {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             ArrayFunctionSignature::ArrayAndElement => {
-                write!(f, "ArrayAndElement")
+                write!(f, "array, element")
             }
             ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {
                 write!(f, "array, element, [index]")
             }
             ArrayFunctionSignature::ElementAndArray => {
-                write!(f, "ElementAndArray")
+                write!(f, "element, array")
             }
             ArrayFunctionSignature::ArrayAndIndex => {
-                write!(f, "ArrayAndIndex")
+                write!(f, "array, index")
             }
             ArrayFunctionSignature::Array => {
-                write!(f, "Array")
+                write!(f, "array")
             }
         }
     }

datafusion/sqllogictest/test_files/array.slt

@@ -1116,7 +1116,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\(ArrayAndIndex\)
+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\)
 select array_element(1, 2);
 
 # array_element with null
@@ -2038,16 +2038,13 @@ select
 ----
 [4] [] [1, , 3, 4] [, , 1]
 
-#TODO: https://github.com/apache/arrow-datafusion/issues/9158
-#query ????
-#select
-#  array_append(arrow_cast(make_array(), 'FixedSizeList(1, Null)'), 4),
-#  array_append(arrow_cast(make_array(), 'FixedSizeList(1, Null)'), null),
-#  array_append(arrow_cast(make_array(1, null, 3), 'FixedSizeList(3, Int64)'), 4),
-#  array_append(arrow_cast(make_array(null, null), 'FixedSizeList(2, Null)'), 1)
-#;
-#----
-#[4] [] [1, , 3, 4] [, , 1]
+query ??
+select
+  array_append(arrow_cast(make_array(1, null, 3), 'FixedSizeList(3, Int64)'), 4),
+  array_append(arrow_cast(make_array(null, null), 'FixedSizeList(2, Int64)'), 1)
+;
+----
+[1, , 3, 4] [, , 1]
 
 # test invalid (non-null)
 query error
@@ -4607,7 +4604,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\)
 select array_dims(1);
 
 # array_dims scalar function
@@ -4621,12 +4618,27 @@ select array_dims(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), array_di
 ----
 [3] [2, 2] [1, 1, 1, 2, 1]
 
+query ???
+select array_dims(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)')), array_dims(arrow_cast(make_array([1, 2], [3, 4]), 'FixedSizeList(2, List(Int64))')), array_dims(arrow_cast(make_array([[[[1], [2]]]]), 'FixedSizeList(1, List(List(List(List(Int64)))))'));
+----
+[3] [2, 2] [1, 1, 1, 2, 1]
+
 # array_dims scalar function #2
 query ??
 select array_dims(array_repeat(array_repeat(array_repeat(2, 3), 2), 1)), array_dims(array_repeat(array_repeat(array_repeat(3, 4), 5), 2));
 ----
 [1, 2, 3] [2, 5, 4]
 
+query ??
+select array_dims(arrow_cast(array_repeat(array_repeat(array_repeat(2, 3), 2), 1), 'LargeList(List(List(Int64)))')), array_dims(arrow_cast(array_repeat(array_repeat(array_repeat(3, 4), 5), 2), 'LargeList(List(List(Int64)))'));
+----
+[1, 2, 3] [2, 5, 4]
+
+query ??
+select array_dims(arrow_cast(array_repeat(array_repeat(array_repeat(2, 3), 2), 1), 'FixedSizeList(1, List(List(Int64)))')), array_dims(arrow_cast(array_repeat(array_repeat(array_repeat(3, 4), 5), 2), 'FixedSizeList(2, List(List(Int64)))'));
+----
+[1, 2, 3] [2, 5, 4]
+
 # array_dims scalar function #3
 query ??
 select array_dims(make_array()), array_dims(make_array(make_array()))
@@ -4649,6 +4661,11 @@ select list_dims(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), list_dims
 ----
 [3] [2, 2] [1, 1, 1, 2, 1]
 
+query ???
+select list_dims(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)')), list_dims(arrow_cast(make_array([1, 2], [3, 4]), 'FixedSizeList(2, List(Int64))')), list_dims(arrow_cast(make_array([[[[1], [2]]]]), 'FixedSizeList(1, List(List(List(List(Int64)))))'));
+----
+[3] [2, 2] [1, 1, 1, 2, 1]
+
 # array_dims with columns
 query ???
 select array_dims(column1), array_dims(column2), array_dims(column3) from arrays;
@@ -4672,6 +4689,17 @@ NULL [3] [4]
 [2, 2] NULL [1]
 [2, 2] [3] NULL
 
+query ???
+select array_dims(column1), array_dims(column2), array_dims(column3) from fixed_size_arrays;
+----
+[2, 2] [3] [5]
+[2, 2] [3] [5]
+[2, 2] [3] [5]
+[2, 2] [3] [5]
+NULL [3] [5]
+[2, 2] NULL [5]
+[2, 2] [3] NULL
+
 
 ## array_ndims (aliases: `list_ndims`)