chore: add tests

apache · Nov 15, 2023 · 16d3738 · 16d3738
1 parent 974a9e7
commit 16d3738
Showing 1 changed file with 26 additions and 20 deletions.
diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
@@ -2161,7 +2161,14 @@ mod tests {
     #[test]
     fn test_return_large_array() {
         let list_array = return_large_array();
-        assert_eq!(list_array.data_type().to_string(), "failed to cast to");
+        let result = as_large_list_array(&list_array);
+        assert_eq!(result.len(), 1);
+        assert_eq!(
+            &[1, 2, 3, 4],
+            as_int64_array(&result.value(0))
+                .expect("failed to cast to primitive array")
+                .values()
+        );
     }
 
     #[test]
@@ -2910,8 +2917,6 @@ mod tests {
     fn test_array_length() {
         // array_length([1, 2, 3, 4]) = 4
         let list_array = return_array();
-        // cast thr list_array to LargeList
-        println!("{:?}", list_array.data_type());
         let arr = array_length(&[list_array.clone()])
             .expect("failed to initialize function array_ndims");
         let result =
@@ -2929,9 +2934,19 @@ mod tests {
         assert_eq!(result, &UInt64Array::from_value(4, 1));
 
         // for LargeList
-        // array_length([1, 2, 3, 4], 2) = 4
-        let array = array_length(&[list_array, Arc::new(Int64Array::from_value(1, 1))])
+        // array_length([1, 2, 3, 4]) = 4
+        let list_array = return_large_array();
+        let arr = array_length(&[list_array.clone()])
             .expect("failed to initialize function array_ndims");
+        let result =
+            as_uint64_array(&arr).expect("failed to initialize function array_ndims");
+
+        assert_eq!(result, &UInt64Array::from_value(4, 1));
+
+        // array_length([1, 2, 3, 4], 1) = 4
+        let array =
+            array_length(&[list_array.clone(), Arc::new(Int64Array::from_value(1, 1))])
+                .expect("failed to initialize function array_ndims");
         let result =
             as_uint64_array(&array).expect("failed to initialize function array_ndims");
 
@@ -3072,23 +3087,14 @@ mod tests {
 
     fn return_large_array() -> ArrayRef {
         // Returns: [1, 2, 3, 4]
-        let capacity = i32::MAX as usize + 10;
-        let args = vec![Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef; capacity];
-
-        println!("args.len() = {}", args.len());
-
-        make_array(&args).expect("failed to initialize function array")
-    }
-
-    fn return_extra_array() -> ArrayRef {
-        // Returns: [11, 12, 13, 14]
         let args = [
-            Arc::new(Int64Array::from(vec![Some(11)])) as ArrayRef,
-            Arc::new(Int64Array::from(vec![Some(12)])) as ArrayRef,
-            Arc::new(Int64Array::from(vec![Some(13)])) as ArrayRef,
-            Arc::new(Int64Array::from(vec![Some(14)])) as ArrayRef,
+            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,
         ];
-        make_array(&args).expect("failed to initialize function array")
+        let data_type = DataType::Int64;
+        array_array::<i64>(&args, data_type).expect("failed to initialize function array")
     }
 
     fn return_nested_array() -> ArrayRef {