feat: support SQL array replacement and removement functions

datafusion/expr/src/built_in_function.rs

@@ -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
@@ -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,
@@ -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;
@@ -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() {
@@ -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())
@@ -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",

datafusion/expr/src/expr_fn.rs

@@ -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,
@@ -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);