add array_and_index function

apache · Jan 30, 2024 · 76cede4 · 76cede4
1 parent 159648a
commit 76cede4
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 11 deletions.
diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
@@ -120,11 +120,7 @@ fn get_valid_types(
             DataType::List(ref field)
             | DataType::LargeList(ref field)
             | DataType::FixedSizeList(ref field, _) => {
-                let elem_type = if elem_base_type.eq(&DataType::Null) {
-                    field.data_type()
-                } else {
-                    elem_type
-                };
+                let elem_type = field.data_type();
                 if is_append {
                     Ok(vec![vec![array_type.clone(), elem_type.clone()]])
                 } else {
@@ -134,6 +130,22 @@ fn get_valid_types(
             _ => 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];
+
+        match array_type {
+            DataType::List(_)
+            | DataType::LargeList(_)
+            | DataType::FixedSizeList(_, _) => {
+                Ok(vec![vec![array_type.clone(), DataType::Int64]])
+            }
+            _ => Ok(vec![vec![]]),
+        }
+    }
     let valid_types = match signature {
         TypeSignature::Variadic(valid_types) => valid_types
             .iter()
@@ -168,10 +180,12 @@ fn get_valid_types(
         TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
         TypeSignature::ArraySignature(ref function_signature) => match function_signature
         {
-            ArrayFunctionSignature::ArrayAndElement
-            | ArrayFunctionSignature::ArrayAndIndex => {
+            ArrayFunctionSignature::ArrayAndElement => {
                 return array_append_or_prepend_valid_types(current_types, true)
             }
+            ArrayFunctionSignature::ArrayAndIndex => {
+                return array_and_index(current_types)
+            }
             ArrayFunctionSignature::ElementAndArray => {
                 return array_append_or_prepend_valid_types(current_types, false)
             }

diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt
@@ -1053,7 +1053,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\(ArrayAndElement\(List<T>, T\)\)
+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
@@ -1114,14 +1114,20 @@ select array_element(arrow_cast(make_array(1, 2, 3, 4, 5), 'FixedSizeList(5, Int
 NULL NULL
 
 # array_element scalar function #4 (with NULL)
-query error
+query IT
 select array_element(make_array(1, 2, 3, 4, 5), NULL), array_element(make_array('h', 'e', 'l', 'l', 'o'), NULL);
+----
+NULL NULL
 
-query error
+query IT
 select array_element(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), NULL), array_element(arrow_cast(make_array('h', 'e', 'l', 'l', 'o'), 'LargeList(Utf8)'), NULL);
+----
+NULL NULL
 
-query error
+query IT
 select array_element(arrow_cast(make_array(1, 2, 3, 4, 5), 'FixedSizeList(5, Int64)'), NULL), array_element(arrow_cast(make_array('h', 'e', 'l', 'l', 'o'), 'FixedSizeList(5, Utf8)'), NULL);
+----
+NULL NULL
 
 # array_element scalar function #5 (with negative index)
 query IT