fix: update tests for nested array

apache · Nov 17, 2023 · d340ff7 · d340ff7
1 parent a617329
commit d340ff7
Showing 1 changed file with 68 additions and 0 deletions.
diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
@@ -171,6 +171,10 @@ fn compute_array_length(
                 value = downcast_arg!(value, ListArray).value(0);
                 current_dimension += 1;
             }
+            DataType::LargeList(..) => {
+                value = downcast_arg!(value, LargeListArray).value(0);
+                current_dimension += 1;
+            }
             _ => return Ok(None),
         }
     }
@@ -3007,6 +3011,43 @@ mod tests {
             as_uint64_array(&arr).expect("failed to initialize function array_length");
 
         assert_eq!(result, &UInt64Array::from(vec![None]));
+
+        // for LargeList
+        let list_array = return_nested_large_array();
+
+        // array_length([[1, 2, 3, 4], [5, 6, 7, 8]]) = 2
+        let arr = array_length(&[list_array.clone()])
+            .expect("failed to initialize function array_length");
+        let result =
+            as_uint64_array(&arr).expect("failed to initialize function array_length");
+
+        assert_eq!(result, &UInt64Array::from_value(2, 1));
+
+        // array_length([[1, 2, 3, 4], [5, 6, 7, 8]], 1) = 2
+        let arr =
+            array_length(&[list_array.clone(), Arc::new(Int64Array::from_value(1, 1))])
+                .expect("failed to initialize function array_length");
+        let result =
+            as_uint64_array(&arr).expect("failed to initialize function array_length");
+
+        assert_eq!(result, &UInt64Array::from_value(2, 1));
+
+        // array_length([[1, 2, 3, 4], [5, 6, 7, 8]], 2) = 4
+        let arr =
+            array_length(&[list_array.clone(), Arc::new(Int64Array::from_value(2, 1))])
+                .expect("failed to initialize function array_length");
+        let result =
+            as_uint64_array(&arr).expect("failed to initialize function array_length");
+
+        assert_eq!(result, &UInt64Array::from_value(4, 1));
+
+        // array_length([[1, 2, 3, 4], [5, 6, 7, 8]], 3) = NULL
+        let arr = array_length(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
+            .expect("failed to initialize function array_length");
+        let result =
+            as_uint64_array(&arr).expect("failed to initialize function array_length");
+
+        assert_eq!(result, &UInt64Array::from(vec![None]));
     }
 
     #[test]
@@ -3135,6 +3176,33 @@ mod tests {
         make_array(&[arr1, arr2]).expect("failed to initialize function array")
     }
 
+    fn return_nested_large_array() -> ArrayRef {
+        // Returns: [[1, 2, 3, 4], [5, 6, 7, 8]]
+        let args = [
+            Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
+            Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
+            Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
+            Arc::new(Int64Array::from(vec![Some(4)])) as ArrayRef,
+        ];
+        let data_type = DataType::Int64;
+        let arr1 = array_array::<i64>(&args, data_type.clone())
+            .expect("failed to initialize function array");
+
+        let args = [
+            Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef,
+            Arc::new(Int64Array::from(vec![Some(6)])) as ArrayRef,
+            Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef,
+            Arc::new(Int64Array::from(vec![Some(8)])) as ArrayRef,
+        ];
+        let arr2 = array_array::<i64>(&args, data_type)
+            .expect("failed to initialize function array");
+        let data_type =
+            DataType::LargeList(Arc::new(Field::new("item", DataType::Int64, true)));
+
+        array_array::<i64>(&[arr1, arr2], data_type)
+            .expect("failed to initialize function array")
+    }
+
     fn return_array_with_nulls() -> ArrayRef {
         // Returns: [1, NULL, 3, NULL]
         let args = [