update test in sqllogictest

apache · Nov 22, 2023 · 5d5d0da · 5d5d0da
1 parent 73f3b06
commit 5d5d0da
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 115 deletions.
diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
@@ -2071,20 +2071,20 @@ impl ScalarValue {
     /// Example
     /// ```
     /// use datafusion_common::ScalarValue;
-    /// use arrow::array::{LargeListArray, Int64Array};
-    /// use arrow::datatypes::{DataType, Int64Type};
+    /// use arrow::array::{LargeListArray, Int32Array};
+    /// use arrow::datatypes::{DataType, Int32Type};
     /// use datafusion_common::cast::as_large_list_array;
     ///
     /// let scalars = vec![
-    ///    ScalarValue::Int64(Some(1)),
-    ///    ScalarValue::Int64(None),
-    ///    ScalarValue::Int64(Some(2))
+    ///    ScalarValue::Int32(Some(1)),
+    ///    ScalarValue::Int32(None),
+    ///    ScalarValue::Int32(Some(2))
     /// ];
     ///
-    /// let array = ScalarValue::new_large_list(&scalars, &DataType::Int64);
+    /// let array = ScalarValue::new_large_list(&scalars, &DataType::Int32);
     /// let result = as_large_list_array(&array).unwrap();
     ///
-    /// let expected = LargeListArray::from_iter_primitive::<Int64Type, _, _>(
+    /// let expected = LargeListArray::from_iter_primitive::<Int32Type, _, _>(
     ///     vec![
     ///        Some(vec![Some(1), None, Some(2)])
     ///     ]);

diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
@@ -2120,19 +2120,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn test_return_large_array() {
-        let list_array = return_large_array();
-        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]
     fn test_array_element() {
         // array_element([1, 2, 3, 4], 1) = 1
@@ -2894,25 +2881,6 @@ mod tests {
             as_uint64_array(&array).expect("failed to initialize function array_ndims");
 
         assert_eq!(result, &UInt64Array::from_value(4, 1));
-
-        // for LargeList
-        // 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");
-
-        assert_eq!(result, &UInt64Array::from_value(4, 1));
     }
 
     #[test]
@@ -2952,43 +2920,6 @@ 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]
@@ -3084,18 +3015,6 @@ mod tests {
         make_array(&args).expect("failed to initialize function array")
     }
 
-    fn return_large_array() -> ArrayRef {
-        // Returns: [1, 2, 3, 4]
-        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;
-        array_array::<i64>(&args, data_type).expect("failed to initialize function array")
-    }
-
     fn return_nested_array() -> ArrayRef {
         // Returns: [[1, 2, 3, 4], [5, 6, 7, 8]]
         let args = [
@@ -3117,33 +3036,6 @@ 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 = [

diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt
@@ -2371,36 +2371,66 @@ select array_length(make_array(1, 2, 3, 4, 5)), array_length(make_array(1, 2, 3)
 ----
 5 3 3
 
+query III
+select array_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)')), array_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), array_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'));
+----
+5 3 3
+
 # array_length scalar function #2
 query III
 select array_length(make_array(1, 2, 3, 4, 5), 1), array_length(make_array(1, 2, 3), 1), array_length(make_array([1, 2], [3, 4], [5, 6]), 1);
 ----
 5 3 3
 
+query III
+select array_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), 1), array_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), 1), array_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'), 1);
+----
+5 3 3
+
 # array_length scalar function #3
 query III
 select array_length(make_array(1, 2, 3, 4, 5), 2), array_length(make_array(1, 2, 3), 2), array_length(make_array([1, 2], [3, 4], [5, 6]), 2);
 ----
 NULL NULL 2
 
+query III
+select array_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), 2), array_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), 2), array_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'), 2);
+----
+NULL NULL 2
+
 # array_length scalar function #4
 query II
 select array_length(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 1), array_length(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 2);
 ----
 3 2
 
+query II
+select array_length(arrow_cast(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 'LargeList(List(List(Int64)))'), 1), array_length(arrow_cast(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 'LargeList(List(List(Int64)))'), 2);
+----
+3 2
+
 # array_length scalar function #5
 query III
 select array_length(make_array()), array_length(make_array(), 1), array_length(make_array(), 2)
 ----
 0 0 NULL
 
+query III
+select array_length(arrow_cast(make_array(), 'LargeList(Null)')), array_length(arrow_cast(make_array(), 'LargeList(Null)'), 1), array_length(arrow_cast(make_array(), 'LargeList(Null)'), 2)
+----
+0 0 NULL
+
 # list_length scalar function #6 (function alias `array_length`)
 query III
 select list_length(make_array(1, 2, 3, 4, 5)), list_length(make_array(1, 2, 3)), list_length(make_array([1, 2], [3, 4], [5, 6]));
 ----
 5 3 3
 
+query III
+select list_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)')), list_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), list_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'));
+----
+5 3 3
+
 # array_length with columns
 query I
 select array_length(column1, column3) from arrays_values;
@@ -2414,6 +2444,18 @@ NULL
 NULL
 NULL
 
+query I
+select array_length(arrow_cast(column1, 'LargeList(Int64)'), column3) from arrays_values;
+----
+10
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+
 # array_length with columns and scalars
 query II
 select array_length(array[array[1, 2], array[3, 4]], column3), array_length(column1, 1) from arrays_values;
@@ -2427,6 +2469,18 @@ NULL 10
 NULL 10
 NULL 10
 
+query II
+select array_length(arrow_cast(array[array[1, 2], array[3, 4]], 'LargeList(List(Int64))'), column3), array_length(arrow_cast(column1, 'LargeList(Int64)'), 1) from arrays_values;
+----
+2 10
+2 10
+NULL 10
+NULL 10
+NULL NULL
+NULL 10
+NULL 10
+NULL 10
+
 ## array_dims (aliases: `list_dims`)
 
 # array dims error