diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 92f9051247aa9..53d12af037984 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -171,6 +171,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), } } @@ -3007,6 +3011,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] @@ -3135,6 +3176,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 = [