Skip to content

Commit

Permalink
fix: update tests for nested array
Browse files Browse the repository at this point in the history
  • Loading branch information
Weijun-H committed Nov 16, 2023
1 parent 16d3738 commit a06fbf8
Showing 1 changed file with 90 additions and 6 deletions.
96 changes: 90 additions & 6 deletions datafusion/physical-expr/src/array_expressions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
}
Expand Down Expand Up @@ -325,7 +329,7 @@ fn array_array<O: OffsetSizeTrait>(
}

let res: Arc<GenericListArray<O>> = 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![];
Expand All @@ -341,11 +345,27 @@ fn array_array<O: OffsetSizeTrait>(
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::<i32>(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::<i64>(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() {
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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::<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 = [
Expand Down

0 comments on commit a06fbf8

Please sign in to comment.