Skip to content
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

Jump to bottom

feat: support FixedSizeList Type Coercion #9108

Merged
merged 21 commits into from
Feb 26, 2024
Merged

feat: support FixedSizeList Type Coercion #9108

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
56e4ddf
support FixedSizeList Type Coercion
Weijun-H Feb 2, 2024
c10b8b2
add allow null type coercion parameter
Weijun-H Feb 2, 2024
92d5e9e
support null column in FixedSizeList
Weijun-H Feb 2, 2024
b61cbef
Add test
Weijun-H Feb 7, 2024
3220da5
Add tests for cardinality with fixed size lists
Weijun-H Feb 8, 2024
b565290
chore
Weijun-H Feb 8, 2024
58b7d09
fix ci
Weijun-H Feb 8, 2024
3d5c899
add comment
Weijun-H Feb 13, 2024
eeed955
Fix array_element function signature
Weijun-H Feb 13, 2024
342d379
Remove unused imports and simplify code
Weijun-H Feb 18, 2024
d1f8744
Fix array function signatures and behavior
Weijun-H Feb 19, 2024
4de5ee2
fix conflict
Weijun-H Feb 19, 2024
13801be
fix conflict
Weijun-H Feb 21, 2024
4b21267
add tests for FixedSizeList
Weijun-H Feb 21, 2024
1d9f35e
remove unreacheable null check
Weijun-H Feb 22, 2024
cb37528
simplify the code
Weijun-H Feb 22, 2024
a7f121e
remove null checking
Weijun-H Feb 25, 2024
0a338b6
reformat output
Weijun-H Feb 26, 2024
9d9527e
simplify code
Weijun-H Feb 26, 2024
9201a36
add tests for array_dims
Weijun-H Feb 26, 2024
ad73b17
Refactor type coercion functions in datafusion/expr module
Weijun-H Feb 26, 2024
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
53 changes: 26 additions & 27 deletions datafusion/expr/src/signature.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,33 +236,15 @@ impl ArrayFunctionSignature {
_ => Ok(vec![vec![]]),
}
}
fn array_and_index(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
if current_types.len() != 2 {
return Ok(vec![vec![]]);
}
let array_type = &current_types[0];
let array_type = array(&[array_type.clone()])?;
if array_type[0].is_empty() {
Ok(vec![vec![]])
} else {
Ok(vec![vec![array_type[0][0].clone(), DataType::Int64]])
}
}
fn array(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
if current_types.len() != 1 {
return Ok(vec![vec![]]);
}

let array_type = &current_types[0];

fn array(array_type: &DataType) -> Option<DataType> {
match array_type {
DataType::List(_)
| DataType::LargeList(_)
| DataType::FixedSizeList(_, _) => {
let array_type = coerced_fixed_size_list_to_list(array_type);
Ok(vec![vec![array_type]])
Some(array_type)
}
_ => Ok(vec![vec![]]),
_ => None,
}
}
match self {
Expand All @@ -272,11 +254,28 @@ impl ArrayFunctionSignature {
ArrayFunctionSignature::ElementAndArray => {
array_append_or_prepend_valid_types(current_types, false)
}
ArrayFunctionSignature::ArrayAndIndex => array_and_index(current_types),
ArrayFunctionSignature::ArrayAndIndex => {
if current_types.len() != 2 {
return Ok(vec![vec![]]);
}
array(&current_types[0]).map_or_else(
|| Ok(vec![vec![]]),
|array_type| Ok(vec![vec![array_type, DataType::Int64]]),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

)
}
ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {
array_element_and_optional_index(current_types)
}
ArrayFunctionSignature::Array => array(current_types),
ArrayFunctionSignature::Array => {
if current_types.len() != 1 {
return Ok(vec![vec![]]);
}

array(&current_types[0]).map_or_else(
|| Ok(vec![vec![]]),
|array_type| Ok(vec![vec![array_type]]),
)
}
}
}
}
Expand All @@ -285,19 +284,19 @@ impl std::fmt::Display for ArrayFunctionSignature {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ArrayFunctionSignature::ArrayAndElement => {
write!(f, "ArrayAndElement")
write!(f, "array, element")
}
ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {
write!(f, "array, element, [index]")

This comment was marked as outdated.

Sign in to view

}
ArrayFunctionSignature::ElementAndArray => {
write!(f, "ElementAndArray")
write!(f, "element, array")
}
ArrayFunctionSignature::ArrayAndIndex => {
write!(f, "ArrayAndIndex")
write!(f, "array, index")
}
ArrayFunctionSignature::Array => {
write!(f, "Array")
write!(f, "array")
}
}
}
Expand Down
52 changes: 40 additions & 12 deletions datafusion/sqllogictest/test_files/array.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1116,7 +1116,7 @@ from arrays_values_without_nulls;
## array_element (aliases: array_extract, list_extract, list_element)

# array_element error
query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_element\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_element\(ArrayAndIndex\)
query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_element\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_element\(array, index\)
select array_element(1, 2);

# array_element with null
Expand Down Expand Up @@ -2038,16 +2038,13 @@ select
----
[4] [] [1, , 3, 4] [, , 1]

#TODO: https://github.com/apache/arrow-datafusion/issues/9158
#query ????
#select
# array_append(arrow_cast(make_array(), 'FixedSizeList(1, Null)'), 4),
# array_append(arrow_cast(make_array(), 'FixedSizeList(1, Null)'), null),
# array_append(arrow_cast(make_array(1, null, 3), 'FixedSizeList(3, Int64)'), 4),
# array_append(arrow_cast(make_array(null, null), 'FixedSizeList(2, Null)'), 1)
#;
#----
#[4] [] [1, , 3, 4] [, , 1]
query ??
select
array_append(arrow_cast(make_array(1, null, 3), 'FixedSizeList(3, Int64)'), 4),
array_append(arrow_cast(make_array(null, null), 'FixedSizeList(2, Int64)'), 1)
;
----
[1, , 3, 4] [, , 1]

# test invalid (non-null)
query error
Expand Down Expand Up @@ -4607,7 +4604,7 @@ NULL 10
## array_dims (aliases: `list_dims`)

# array dims error
query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_dims\(Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_dims\(Array\)
query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_dims\(Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_dims\(array\)
select array_dims(1);

# array_dims scalar function
Expand All @@ -4621,12 +4618,27 @@ select array_dims(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), array_di
----
[3] [2, 2] [1, 1, 1, 2, 1]

query ???
select array_dims(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)')), array_dims(arrow_cast(make_array([1, 2], [3, 4]), 'FixedSizeList(2, List(Int64))')), array_dims(arrow_cast(make_array([[[[1], [2]]]]), 'FixedSizeList(1, List(List(List(List(Int64)))))'));
----
[3] [2, 2] [1, 1, 1, 2, 1]

# array_dims scalar function #2
query ??
select array_dims(array_repeat(array_repeat(array_repeat(2, 3), 2), 1)), array_dims(array_repeat(array_repeat(array_repeat(3, 4), 5), 2));
----
[1, 2, 3] [2, 5, 4]

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

query ??
select array_dims(arrow_cast(array_repeat(array_repeat(array_repeat(2, 3), 2), 1), 'FixedSizeList(1, List(List(Int64)))')), array_dims(arrow_cast(array_repeat(array_repeat(array_repeat(3, 4), 5), 2), 'FixedSizeList(2, List(List(Int64)))'));
----
[1, 2, 3] [2, 5, 4]

# array_dims scalar function #3
query ??
select array_dims(make_array()), array_dims(make_array(make_array()))
Expand All @@ -4649,6 +4661,11 @@ select list_dims(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), list_dims
----
[3] [2, 2] [1, 1, 1, 2, 1]

query ???
select list_dims(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)')), list_dims(arrow_cast(make_array([1, 2], [3, 4]), 'FixedSizeList(2, List(Int64))')), list_dims(arrow_cast(make_array([[[[1], [2]]]]), 'FixedSizeList(1, List(List(List(List(Int64)))))'));
----
[3] [2, 2] [1, 1, 1, 2, 1]

# array_dims with columns
query ???
select array_dims(column1), array_dims(column2), array_dims(column3) from arrays;
Expand All @@ -4672,6 +4689,17 @@ NULL [3] [4]
[2, 2] NULL [1]
[2, 2] [3] NULL

query ???
select array_dims(column1), array_dims(column2), array_dims(column3) from fixed_size_arrays;
----
[2, 2] [3] [5]
[2, 2] [3] [5]
[2, 2] [3] [5]
[2, 2] [3] [5]
NULL [3] [5]
[2, 2] NULL [5]
[2, 2] [3] NULL


## array_ndims (aliases: `list_ndims`)

Expand Down
Loading