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 SQL array replacement and removement functions #7057

Merged
merged 13 commits into from
Jul 27, 2023
Merged

feat: support SQL array replacement and removement functions #7057

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
41442f3
feat: sqllogictests for array_remove, array_replace, array_removes an…
izveigor Jul 20, 2023
2b846b1
test: unit tests for array_replace, array_replaces, array_remove, arr…
izveigor Jul 20, 2023
909a07f
feat: array_replace, array_replaces, array_remove, array_removes
izveigor Jul 22, 2023
1e4150c
docs: array_replace, array_replaces, array_remove, array_removes
izveigor Jul 22, 2023
97c164d
Merge branch main
izveigor Jul 22, 2023
e6bc578
feat: proto
izveigor Jul 22, 2023
9428b78
fix: use prettier
izveigor Jul 22, 2023
1d045ef
fix: cargo fmt
izveigor Jul 22, 2023
b944d6f
feat: optional argument for array_replace and array_remove
izveigor Jul 24, 2023
c7021af
fix: proto
izveigor Jul 24, 2023
7f86621
fix: expr_fn tests
izveigor Jul 24, 2023
6b13efa
Merge branch 'main' into replace_remove_functions
izveigor Jul 24, 2023
be02919
feat: array_replace_n and array_remove_n
izveigor Jul 25, 2023
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
386 changes: 378 additions & 8 deletions datafusion/core/tests/sqllogictests/test_files/array.slt
View file
Open in desktop

Large diffs are not rendered by default.

73 changes: 34 additions & 39 deletions datafusion/expr/src/built_in_function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,8 +141,16 @@ pub enum BuiltinScalarFunction {
ArrayPrepend,
/// array_remove
ArrayRemove,
/// array_remove_n
ArrayRemoveN,
/// array_remove_all
ArrayRemoveAll,
/// array_replace
ArrayReplace,
/// array_replace_n
ArrayReplaceN,
/// array_replace_all
ArrayReplaceAll,
/// array_to_string
ArrayToString,
/// cardinality
Expand Down Expand Up @@ -345,7 +353,11 @@ impl BuiltinScalarFunction {
BuiltinScalarFunction::ArrayPositions => Volatility::Immutable,
BuiltinScalarFunction::ArrayPrepend => Volatility::Immutable,
BuiltinScalarFunction::ArrayRemove => Volatility::Immutable,
BuiltinScalarFunction::ArrayRemoveN => Volatility::Immutable,
BuiltinScalarFunction::ArrayRemoveAll => Volatility::Immutable,
BuiltinScalarFunction::ArrayReplace => Volatility::Immutable,
BuiltinScalarFunction::ArrayReplaceN => Volatility::Immutable,
BuiltinScalarFunction::ArrayReplaceAll => Volatility::Immutable,
BuiltinScalarFunction::ArrayToString => Volatility::Immutable,
BuiltinScalarFunction::Cardinality => Volatility::Immutable,
BuiltinScalarFunction::MakeArray => Volatility::Immutable,
Expand Down Expand Up @@ -473,16 +485,7 @@ impl BuiltinScalarFunction {
// the return type of the built in function.
// Some built-in functions' return type depends on the incoming type.
match self {
BuiltinScalarFunction::ArrayAppend => match &input_expr_types[0] {
List(_) => Ok(List(Arc::new(Field::new(
"item",
input_expr_types[1].clone(),
true,
)))),
_ => Err(DataFusionError::Internal(format!(
"The {self} function can only accept list as the first argument"
))),
},
BuiltinScalarFunction::ArrayAppend => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayConcat => {
let mut expr_type = Null;
let mut max_dims = 0;
Expand Down Expand Up @@ -524,31 +527,13 @@ impl BuiltinScalarFunction {
BuiltinScalarFunction::ArrayPositions => {
Ok(List(Arc::new(Field::new("item", UInt64, true))))
}
BuiltinScalarFunction::ArrayPrepend => Ok(List(Arc::new(Field::new(
"item",
input_expr_types[0].clone(),
true,
)))),
BuiltinScalarFunction::ArrayRemove => match &input_expr_types[0] {
List(field) => Ok(List(Arc::new(Field::new(
"item",
field.data_type().clone(),
true,
)))),
_ => Err(DataFusionError::Internal(format!(
"The {self} function can only accept list as the first argument"
))),
},
BuiltinScalarFunction::ArrayReplace => match &input_expr_types[0] {
List(field) => Ok(List(Arc::new(Field::new(
"item",
field.data_type().clone(),
true,
)))),
_ => Err(DataFusionError::Internal(format!(
"The {self} function can only accept list as the first argument"
))),
},
BuiltinScalarFunction::ArrayPrepend => Ok(input_expr_types[1].clone()),
BuiltinScalarFunction::ArrayRemove => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayRemoveN => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayRemoveAll => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayReplace => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayReplaceN => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayReplaceAll => Ok(input_expr_types[0].clone()),
BuiltinScalarFunction::ArrayToString => Ok(Utf8),
BuiltinScalarFunction::Cardinality => Ok(UInt64),
BuiltinScalarFunction::MakeArray => match input_expr_types.len() {
Expand Down Expand Up @@ -831,8 +816,12 @@ impl BuiltinScalarFunction {
BuiltinScalarFunction::ArrayPositions => Signature::any(2, self.volatility()),
BuiltinScalarFunction::ArrayPrepend => Signature::any(2, self.volatility()),
BuiltinScalarFunction::ArrayRemove => Signature::any(2, self.volatility()),
BuiltinScalarFunction::ArrayReplace => {
Signature::variadic_any(self.volatility())
BuiltinScalarFunction::ArrayRemoveN => Signature::any(3, self.volatility()),
BuiltinScalarFunction::ArrayRemoveAll => Signature::any(2, self.volatility()),
BuiltinScalarFunction::ArrayReplace => Signature::any(3, self.volatility()),
BuiltinScalarFunction::ArrayReplaceN => Signature::any(4, self.volatility()),
BuiltinScalarFunction::ArrayReplaceAll => {
Signature::any(3, self.volatility())
}
BuiltinScalarFunction::ArrayToString => {
Signature::variadic_any(self.volatility())
Expand Down Expand Up @@ -1310,8 +1299,14 @@ fn aliases(func: &BuiltinScalarFunction) -> &'static [&'static str] {
"array_push_front",
"list_push_front",
],
BuiltinScalarFunction::ArrayRemove => &["array_remove"],
BuiltinScalarFunction::ArrayReplace => &["array_replace"],
BuiltinScalarFunction::ArrayRemove => &["array_remove", "list_remove"],
BuiltinScalarFunction::ArrayRemoveN => &["array_remove_n", "list_remove_n"],
BuiltinScalarFunction::ArrayRemoveAll => &["array_remove_all", "list_remove_all"],
BuiltinScalarFunction::ArrayReplace => &["array_replace", "list_replace"],
BuiltinScalarFunction::ArrayReplaceN => &["array_replace_n", "list_replace_n"],
BuiltinScalarFunction::ArrayReplaceAll => {
&["array_replace_all", "list_replace_all"]
}
BuiltinScalarFunction::ArrayToString => &[
"array_to_string",
"list_to_string",
Expand Down
32 changes: 30 additions & 2 deletions datafusion/expr/src/expr_fn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -599,13 +599,37 @@ scalar_expr!(
ArrayRemove,
array_remove,
array element,
"removes all elements equal to the given value from the array."
"removes the first element from the array equal to the given value."
);
scalar_expr!(
ArrayRemoveN,
array_remove_n,
array element max,
"removes the first `max` elements from the array equal to the given value."
);
scalar_expr!(
ArrayRemoveAll,
array_remove_all,
array element,
"removes all elements from the array equal to the given value."
);
scalar_expr!(
ArrayReplace,
array_replace,
array from to,
"replaces a specified element with another specified element."
"replaces the first occurrence of the specified element with another specified element."
);
scalar_expr!(
ArrayReplaceN,
array_replace_n,
array from to max,
"replaces the first `max` occurrences of the specified element with another specified element."
);
scalar_expr!(
ArrayReplaceAll,
array_replace_all,
array from to,
"replaces all occurrences of the specified element with another specified element."
);
scalar_expr!(
ArrayToString,
Expand Down Expand Up @@ -1026,7 +1050,11 @@ mod test {
test_scalar_expr!(ArrayPositions, array_positions, array, element);
test_scalar_expr!(ArrayPrepend, array_prepend, array, element);
test_scalar_expr!(ArrayRemove, array_remove, array, element);
test_scalar_expr!(ArrayRemoveN, array_remove_n, array, element, max);
test_scalar_expr!(ArrayRemoveAll, array_remove_all, array, element);
test_scalar_expr!(ArrayReplace, array_replace, array, from, to);
test_scalar_expr!(ArrayReplaceN, array_replace_n, array, from, to, max);
test_scalar_expr!(ArrayReplaceAll, array_replace_all, array, from, to);
test_scalar_expr!(ArrayToString, array_to_string, array, delimiter);
test_unary_scalar_expr!(Cardinality, cardinality);
test_nary_scalar_expr!(MakeArray, array, input);
Expand Down
Loading