Skip to content

feat: support LargeList in array_positions #8571

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 18, 2023
Merged

feat: support LargeList in array_positions #8571

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions datafusion/physical-expr/src/array_expressions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1284,12 +1284,23 @@ fn general_position<OffsetSize: OffsetSizeTrait>(

/// Array_positions SQL function
pub fn array_positions(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr = as_list_array(&args[0])?;
let element = &args[1];

check_datatypes("array_positions", &[arr.values(), element])?;

general_positions::<i32>(arr, element)
match &args[0].data_type() {
DataType::List(_) => {
let arr = as_list_array(&args[0])?;
check_datatypes("array_positions", &[arr.values(), element])?;
general_positions::<i32>(arr, element)
}
DataType::LargeList(_) => {
let arr = as_large_list_array(&args[0])?;
check_datatypes("array_positions", &[arr.values(), element])?;
general_positions::<i64>(arr, element)
}
array_type => {
not_impl_err!("array_positions does not support type '{array_type:?}'.")
}
}
}

fn general_positions<OffsetSize: OffsetSizeTrait>(
Expand Down
43 changes: 43 additions & 0 deletions datafusion/sqllogictest/test_files/array.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1809,18 +1809,33 @@ select array_positions(['h', 'e', 'l', 'l', 'o'], 'l'), array_positions([1, 2, 3
----
[3, 4] [5] [1, 2, 3]

query ???
select array_positions(arrow_cast(['h', 'e', 'l', 'l', 'o'], 'LargeList(Utf8)'), 'l'), array_positions(arrow_cast([1, 2, 3, 4, 5], 'LargeList(Int64)'), 5), array_positions(arrow_cast([1, 1, 1], 'LargeList(Int64)'), 1);
----
[3, 4] [5] [1, 2, 3]

# array_positions scalar function #2 (element is list)
query ?
select array_positions(make_array([1, 2, 3], [2, 1, 3], [1, 5, 6], [2, 1, 3], [4, 5, 6]), [2, 1, 3]);
----
[2, 4]

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

# list_positions scalar function #3 (function alias `array_positions`)
query ???
select list_positions(['h', 'e', 'l', 'l', 'o'], 'l'), list_positions([1, 2, 3, 4, 5], 5), list_positions([1, 1, 1], 1);
----
[3, 4] [5] [1, 2, 3]

query ???
select list_positions(arrow_cast(['h', 'e', 'l', 'l', 'o'], 'LargeList(Utf8)'), 'l'), list_positions(arrow_cast([1, 2, 3, 4, 5], 'LargeList(Int64)'), 5), list_positions(arrow_cast([1, 1, 1], 'LargeList(Int64)'), 1);
----
[3, 4] [5] [1, 2, 3]

# array_positions with columns #1
query ?
select array_positions(column1, column2) from arrays_values_without_nulls;
Expand All @@ -1830,13 +1845,27 @@ select array_positions(column1, column2) from arrays_values_without_nulls;
[3]
[4]

query ?
select array_positions(arrow_cast(column1, 'LargeList(Int64)'), column2) from arrays_values_without_nulls;
----
[1]
[2]
[3]
[4]

# array_positions with columns #2 (element is list)
query ?
select array_positions(column1, column2) from nested_arrays;
----
[3]
[2, 5]

query ?
select array_positions(arrow_cast(column1, 'LargeList(List(Int64))'), column2) from nested_arrays;
----
[3]
[2, 5]

# array_positions with columns and scalars #1
query ??
select array_positions(column1, 4), array_positions(array[1, 2, 23, 13, 33, 45], column2) from arrays_values_without_nulls;
Expand All @@ -1846,13 +1875,27 @@ select array_positions(column1, 4), array_positions(array[1, 2, 23, 13, 33, 45],
[] [3]
[] []

query ??
select array_positions(arrow_cast(column1, 'LargeList(Int64)'), 4), array_positions(array[1, 2, 23, 13, 33, 45], column2) from arrays_values_without_nulls;
----
[4] [1]
[] []
[] [3]
[] []

# array_positions with columns and scalars #2 (element is list)
query ??
select array_positions(column1, make_array(4, 5, 6)), array_positions(make_array([1, 2, 3], [11, 12, 13], [4, 5, 6]), column2) from nested_arrays;
----
[6] []
[1] []

query ??
select array_positions(arrow_cast(column1, 'LargeList(List(Int64))'), make_array(4, 5, 6)), array_positions(arrow_cast(make_array([1, 2, 3], [11, 12, 13], [4, 5, 6]), 'LargeList(List(Int64))'), column2) from nested_arrays;
----
[6] []
[1] []

## array_replace (aliases: `list_replace`)

# array_replace scalar function #1
Expand Down