Skip to content

Commit

Permalink
update test in sqllogictest
Browse files Browse the repository at this point in the history
  • Loading branch information
Weijun-H committed Nov 22, 2023
1 parent 73f3b06 commit 6839aaa
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 108 deletions.
108 changes: 0 additions & 108 deletions datafusion/physical-expr/src/array_expressions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2120,19 +2120,6 @@ mod tests {
);
}

#[test]
fn test_return_large_array() {
let list_array = return_large_array();
let result = as_large_list_array(&list_array);
assert_eq!(result.len(), 1);
assert_eq!(
&[1, 2, 3, 4],
as_int64_array(&result.value(0))
.expect("failed to cast to primitive array")
.values()
);
}

#[test]
fn test_array_element() {
// array_element([1, 2, 3, 4], 1) = 1
Expand Down Expand Up @@ -2894,25 +2881,6 @@ mod tests {
as_uint64_array(&array).expect("failed to initialize function array_ndims");

assert_eq!(result, &UInt64Array::from_value(4, 1));

// for LargeList
// array_length([1, 2, 3, 4]) = 4
let list_array = return_large_array();
let arr = array_length(&[list_array.clone()])
.expect("failed to initialize function array_ndims");
let result =
as_uint64_array(&arr).expect("failed to initialize function array_ndims");

assert_eq!(result, &UInt64Array::from_value(4, 1));

// array_length([1, 2, 3, 4], 1) = 4
let array =
array_length(&[list_array.clone(), Arc::new(Int64Array::from_value(1, 1))])
.expect("failed to initialize function array_ndims");
let result =
as_uint64_array(&array).expect("failed to initialize function array_ndims");

assert_eq!(result, &UInt64Array::from_value(4, 1));
}

#[test]
Expand Down Expand Up @@ -2952,43 +2920,6 @@ 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 @@ -3084,18 +3015,6 @@ mod tests {
make_array(&args).expect("failed to initialize function array")
}

fn return_large_array() -> ArrayRef {
// Returns: [1, 2, 3, 4]
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;
array_array::<i64>(&args, data_type).expect("failed to initialize function array")
}

fn return_nested_array() -> ArrayRef {
// Returns: [[1, 2, 3, 4], [5, 6, 7, 8]]
let args = [
Expand All @@ -3117,33 +3036,6 @@ 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
54 changes: 54 additions & 0 deletions datafusion/sqllogictest/test_files/array.slt
Original file line number Diff line number Diff line change
Expand Up @@ -2371,36 +2371,66 @@ select array_length(make_array(1, 2, 3, 4, 5)), array_length(make_array(1, 2, 3)
----
5 3 3

query III
select array_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)')), array_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), array_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'));
----
5 3 3

# array_length scalar function #2
query III
select array_length(make_array(1, 2, 3, 4, 5), 1), array_length(make_array(1, 2, 3), 1), array_length(make_array([1, 2], [3, 4], [5, 6]), 1);
----
5 3 3

query III
select array_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), 1), array_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), 1), array_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'), 1);
----
5 3 3

# array_length scalar function #3
query III
select array_length(make_array(1, 2, 3, 4, 5), 2), array_length(make_array(1, 2, 3), 2), array_length(make_array([1, 2], [3, 4], [5, 6]), 2);
----
NULL NULL 2

query III
select array_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), 2), array_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), 2), array_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'), 2);
----
NULL NULL 2

# array_length scalar function #4
query II
select array_length(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 1), array_length(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 2);
----
3 2

query II
select array_length(arrow_cast(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 'LargeList(List(List(Int64)))'), 1), array_length(arrow_cast(array_repeat(array_repeat(array_repeat(3, 5), 2), 3), 'LargeList(List(List(Int64)))'), 2);
----
3 2

# array_length scalar function #5
query III
select array_length(make_array()), array_length(make_array(), 1), array_length(make_array(), 2)
----
0 0 NULL

query III
select array_length(arrow_cast(make_array(), 'LargeList(Null)')), array_length(arrow_cast(make_array(), 'LargeList(Null)'), 1), array_length(arrow_cast(make_array(), 'LargeList(Null)'), 2)
----
0 0 NULL

# list_length scalar function #6 (function alias `array_length`)
query III
select list_length(make_array(1, 2, 3, 4, 5)), list_length(make_array(1, 2, 3)), list_length(make_array([1, 2], [3, 4], [5, 6]));
----
5 3 3

query III
select list_length(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)')), list_length(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), list_length(arrow_cast(make_array([1, 2], [3, 4], [5, 6]), 'LargeList(List(Int64))'));
----
5 3 3

# array_length with columns
query I
select array_length(column1, column3) from arrays_values;
Expand All @@ -2414,6 +2444,18 @@ NULL
NULL
NULL

query I
select array_length(arrow_cast(column1, 'LargeList(Int64)'), column3) from arrays_values;
----
10
NULL
NULL
NULL
NULL
NULL
NULL
NULL

# array_length with columns and scalars
query II
select array_length(array[array[1, 2], array[3, 4]], column3), array_length(column1, 1) from arrays_values;
Expand All @@ -2427,6 +2469,18 @@ NULL 10
NULL 10
NULL 10

query II
select array_length(arrow_cast(array[array[1, 2], array[3, 4]], 'LargeList(List(Int64))'), column3), array_length(arrow_cast(column1, 'LargeList(Int64)'), 1) from arrays_values;
----
2 10
2 10
NULL 10
NULL 10
NULL NULL
NULL 10
NULL 10
NULL 10

## array_dims (aliases: `list_dims`)

# array dims error
Expand Down

0 comments on commit 6839aaa

Please sign in to comment.