Port ArrayRemove, ArrayRemoveN, ArrayRemoveAll to `functions-ar…

…ray` subcrate (#9656) * Issue-9635 - Port ArrayRemove, ArrayRemoveN, ArrayRemoveAll to functions-array subcrate * Issue-9635 - Address review comment
apache · Mar 18, 2024 · c7a72e6 · c7a72e6
1 parent cab5d0f
commit c7a72e6
Show file tree

Hide file tree

Showing 13 changed files with 456 additions and 239 deletions.
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
@@ -103,12 +103,6 @@ pub enum BuiltinScalarFunction {
     Cot,
 
     // array functions
-    /// array_remove
-    ArrayRemove,
-    /// array_remove_n
-    ArrayRemoveN,
-    /// array_remove_all
-    ArrayRemoveAll,
     /// array_replace
     ArrayReplace,
     /// array_replace_n
@@ -268,9 +262,6 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::Cbrt => Volatility::Immutable,
             BuiltinScalarFunction::Cot => Volatility::Immutable,
             BuiltinScalarFunction::Trunc => 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,
@@ -331,9 +322,6 @@ impl BuiltinScalarFunction {
         // the return type of the built in function.
         // Some built-in functions' return type depends on the incoming type.
         match self {
-            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()),
@@ -489,13 +477,6 @@ impl BuiltinScalarFunction {
 
         // for now, the list is small, as we do not have many built-in functions.
         match self {
-            BuiltinScalarFunction::ArrayRemove => {
-                Signature::array_and_element(self.volatility())
-            }
-            BuiltinScalarFunction::ArrayRemoveN => Signature::any(3, self.volatility()),
-            BuiltinScalarFunction::ArrayRemoveAll => {
-                Signature::array_and_element(self.volatility())
-            }
             BuiltinScalarFunction::ArrayReplace => Signature::any(3, self.volatility()),
             BuiltinScalarFunction::ArrayReplaceN => Signature::any(4, self.volatility()),
             BuiltinScalarFunction::ArrayReplaceAll => {
@@ -800,11 +781,6 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::FindInSet => &["find_in_set"],
 
             // hashing functions
-            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"]

diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
@@ -584,24 +584,6 @@ scalar_expr!(
 scalar_expr!(Uuid, uuid, , "returns uuid v4 as a string value");
 scalar_expr!(Log, log, base x, "logarithm of a `x` for a particular `base`");
 
-scalar_expr!(
-    ArrayRemove,
-    array_remove,
-    array element,
-    "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,
@@ -1164,9 +1146,6 @@ mod test {
         test_scalar_expr!(Trim, trim, string);
         test_scalar_expr!(Upper, upper, string);
 
-        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);

diff --git a/datafusion/functions-array/src/lib.rs b/datafusion/functions-array/src/lib.rs
@@ -35,6 +35,7 @@ mod except;
 mod extract;
 mod kernels;
 mod position;
+mod remove;
 mod rewrite;
 mod set_ops;
 mod udf;
@@ -62,6 +63,9 @@ pub mod expr_fn {
     pub use super::extract::array_slice;
     pub use super::position::array_position;
     pub use super::position::array_positions;
+    pub use super::remove::array_remove;
+    pub use super::remove::array_remove_all;
+    pub use super::remove::array_remove_n;
     pub use super::set_ops::array_distinct;
     pub use super::set_ops::array_intersect;
     pub use super::set_ops::array_union;
@@ -115,6 +119,9 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
         set_ops::array_union_udf(),
         position::array_position_udf(),
         position::array_positions_udf(),
+        remove::array_remove_udf(),
+        remove::array_remove_n_udf(),
+        remove::array_remove_all_udf(),
     ];
     functions.into_iter().try_for_each(|udf| {
         let existing_udf = registry.register_udf(udf)?;