From a06fbf855d69b564136b1547915fe724b0ff1914 Mon Sep 17 00:00:00 2001 From: Weijun Huang Date: Thu, 16 Nov 2023 11:02:13 +0100 Subject: [PATCH] fix: update tests for nested array --- .../physical-expr/src/array_expressions.rs | 96 +++++++++++++++++-- 1 file changed, 90 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 3804ee65e4126..ce1d688862516 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -166,6 +166,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), } } @@ -325,7 +329,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 +345,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 +3010,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 +3175,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::(&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"); + let data_type = + DataType::LargeList(Arc::new(Field::new("item", DataType::Int64, true))); + + array_array::(&[arr1, arr2], data_type) + .expect("failed to initialize function array") + } + fn return_array_with_nulls() -> ArrayRef { // Returns: [1, NULL, 3, NULL] let args = [