diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 3804ee65e4126..f2eed5d6b2173 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -325,7 +325,7 @@ fn array_array( } let res: Arc> = match data_type { - DataType::List(..) => { + DataType::List(..) | DataType::LargeList(_) => { let row_count = args[0].len(); let column_count = args.len(); let mut list_arrays = vec![]; @@ -341,11 +341,27 @@ fn array_array( array_lengths.push(0); valid.append(false); } else { - let list_arr = as_list_array(arg)?; - let arr = list_arr.value(index); - array_lengths.push(arr.len()); - arrays.push(arr); - valid.append(true); + match arg.data_type() { + DataType::List(_) => { + let list_arr = as_generic_list_array::(arg)?; + let arr = list_arr.value(index); + array_lengths.push(arr.len()); + arrays.push(arr); + valid.append(true); + } + DataType::LargeList(_) => { + let list_arr = as_generic_list_array::(arg)?; + let arr = list_arr.value(index); + array_lengths.push(arr.len()); + arrays.push(arr); + valid.append(true); + } + _ => { + return not_impl_err!( + "Array is not implemented for type '{data_type:?}'." + ) + } + } } } if arrays.is_empty() { @@ -2990,6 +3006,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] @@ -3118,6 +3171,30 @@ 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::(&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::(&args, data_type) + .expect("failed to initialize function array"); + + make_array(&[arr1, arr2]).expect("failed to initialize function array") + } + fn return_array_with_nulls() -> ArrayRef { // Returns: [1, NULL, 3, NULL] let args = [