support LargeList in array_ndims (#8716)

Author: Weijun-H

datafusion/common/src/utils.rs

@@ -469,10 +469,11 @@ pub fn coerced_type_with_base_type_only(
 
 /// Compute the number of dimensions in a list data type.
 pub fn list_ndims(data_type: &DataType) -> u64 {
-    if let DataType::List(field) = data_type {
-        1 + list_ndims(field.data_type())
-    } else {
-        0
+    match data_type {
+        DataType::List(field) | DataType::LargeList(field) => {
+            1 + list_ndims(field.data_type())
+        }
+        _ => 0,
     }
 }
 

datafusion/physical-expr/src/array_expressions.rs

@@ -2250,11 +2250,13 @@ pub fn array_ndims(args: &[ArrayRef]) -> Result<ArrayRef> {
         return exec_err!("array_ndims needs one argument");
     }
 
-    if let Some(list_array) = args[0].as_list_opt::<i32>() {
-        let ndims = datafusion_common::utils::list_ndims(list_array.data_type());
+    fn general_list_ndims<O: OffsetSizeTrait>(
+        array: &GenericListArray<O>,
+    ) -> Result<ArrayRef> {
+        let mut data = Vec::new();
+        let ndims = datafusion_common::utils::list_ndims(array.data_type());
 
-        let mut data = vec![];
-        for arr in list_array.iter() {
+        for arr in array.iter() {
             if arr.is_some() {
                 data.push(Some(ndims))
             } else {
@@ -2263,8 +2265,18 @@ pub fn array_ndims(args: &[ArrayRef]) -> Result<ArrayRef> {
         }
 
         Ok(Arc::new(UInt64Array::from(data)) as ArrayRef)
-    } else {
-        Ok(Arc::new(UInt64Array::from(vec![0; args[0].len()])) as ArrayRef)
+    }
+
+    match args[0].data_type() {
+        DataType::List(_) => {
+            let array = as_list_array(&args[0])?;
+            general_list_ndims::<i32>(array)
+        }
+        DataType::LargeList(_) => {
+            let array = as_large_list_array(&args[0])?;
+            general_list_ndims::<i64>(array)
+        }
+        _ => Ok(Arc::new(UInt64Array::from(vec![0; args[0].len()])) as ArrayRef),
     }
 }
 

datafusion/sqllogictest/test_files/array.slt

@@ -3504,7 +3504,7 @@ NULL [3] [4]
 # array_ndims scalar function #1
 
 query III
-select 
+select
   array_ndims(1),
   array_ndims(null),
   array_ndims([2, 3]);
@@ -3520,8 +3520,17 @@ AS VALUES
   (3, [6], [[9]], [[[[[10]]]]])
 ;
 
+statement ok
+CREATE TABLE large_array_ndims_table
+AS SELECT
+  column1,
+  arrow_cast(column2, 'LargeList(Int64)') as column2,
+  arrow_cast(column3, 'LargeList(List(Int64))') as column3,
+  arrow_cast(column4, 'LargeList(List(List(List(List(Int64)))))') as column4
+FROM array_ndims_table;
+
 query IIII
-select 
+select
   array_ndims(column1),
   array_ndims(column2),
   array_ndims(column3),
@@ -3533,9 +3542,25 @@ from array_ndims_table;
 0 1 2 5
 0 1 2 5
 
+query IIII
+select
+  array_ndims(column1),
+  array_ndims(column2),
+  array_ndims(column3),
+  array_ndims(column4)
+from large_array_ndims_table;
+----
+0 1 2 5
+0 1 2 5
+0 1 2 5
+0 1 2 5
+
 statement ok
 drop table array_ndims_table;
 
+statement ok
+drop table large_array_ndims_table
+
 query I
 select array_ndims(arrow_cast([null], 'List(List(List(Int64)))'));
 ----
@@ -3553,14 +3578,29 @@ select array_ndims(make_array()), array_ndims(make_array(make_array()))
 ----
 1 2
 
+query II
+select array_ndims(arrow_cast(make_array(), 'LargeList(Null)')), array_ndims(arrow_cast(make_array(make_array()), 'LargeList(List(Null))'))
+----
+1 2
+
 # list_ndims scalar function #4 (function alias `array_ndims`)
 query III
 select list_ndims(make_array(1, 2, 3)), list_ndims(make_array([1, 2], [3, 4])), list_ndims(make_array([[[[1], [2]]]]));
 ----
 1 2 5
 
+query III
+select list_ndims(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), list_ndims(arrow_cast(make_array([1, 2], [3, 4]), 'LargeList(List(Int64))')), list_ndims(arrow_cast(make_array([[[[1], [2]]]]), 'LargeList(List(List(List(List(Int64)))))'));
+----
+1 2 5
+
 query II
-select array_ndims(make_array()), array_ndims(make_array(make_array()))
+select list_ndims(make_array()), list_ndims(make_array(make_array()))
+----
+1 2
+
+query II
+select list_ndims(arrow_cast(make_array(), 'LargeList(Null)')), list_ndims(arrow_cast(make_array(make_array()), 'LargeList(List(Null))'))
 ----
 1 2
 
@@ -3576,6 +3616,17 @@ NULL 1 1
 2 NULL 1
 2 1 NULL
 
+query III
+select array_ndims(column1), array_ndims(column2), array_ndims(column3) from large_arrays;
+----
+2 1 1
+2 1 1
+2 1 1
+2 1 1
+NULL 1 1
+2 NULL 1
+2 1 NULL
+
 ## array_has/array_has_all/array_has_any
 
 query BBBBBBBBBBBB