Convert Binary Operator StringConcat to Function for array_concat, array_append and array_prepend

datafusion/expr/src/type_coercion/binary.rs

@@ -667,8 +667,6 @@ fn string_concat_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<Da
         (LargeUtf8, from_type) | (from_type, LargeUtf8) => {
             string_concat_internal_coercion(from_type, &LargeUtf8)
         }
-        // TODO: cast between array elements (#6558)
-        (List(_), from_type) | (from_type, List(_)) => Some(from_type.to_owned()),
         _ => None,
     })
 }

datafusion/optimizer/src/analyzer/mod.rs

@@ -17,6 +17,7 @@
 
 pub mod count_wildcard_rule;
 pub mod inline_table_scan;
+pub mod rewrite_expr;
 pub mod subquery;
 pub mod type_coercion;
 
@@ -37,6 +38,8 @@ use log::debug;
 use std::sync::Arc;
 use std::time::Instant;
 
+use self::rewrite_expr::OperatorToFunction;
+
 /// [`AnalyzerRule`]s transform [`LogicalPlan`]s in some way to make
 /// the plan valid prior to the rest of the DataFusion optimization process.
 ///
@@ -72,6 +75,9 @@ impl Analyzer {
     pub fn new() -> Self {
         let rules: Vec<Arc<dyn AnalyzerRule + Send + Sync>> = vec![
             Arc::new(InlineTableScan::new()),
+            // OperatorToFunction should be run before TypeCoercion, since it rewrite based on the argument types (List or Scalar),
+            // and TypeCoercion may cast the argument types from Scalar to List.
+            Arc::new(OperatorToFunction::new()),
             Arc::new(TypeCoercion::new()),
             Arc::new(CountWildcardRule::new()),
         ];

datafusion/optimizer/src/analyzer/rewrite_expr.rs

@@ -0,0 +1,232 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Optimizer rule for expression rewrite
+
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::TreeNode;
+use datafusion_common::tree_node::TreeNodeRewriter;
+use datafusion_common::Result;
+use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::BuiltinScalarFunction;
+use datafusion_expr::Operator;
+use datafusion_expr::Projection;
+use datafusion_expr::ScalarFunctionDefinition;
+use datafusion_expr::{BinaryExpr, Expr, LogicalPlan};
+
+use super::AnalyzerRule;
+
+#[derive(Default)]
+pub struct OperatorToFunction {}
+
+impl OperatorToFunction {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl AnalyzerRule for OperatorToFunction {
+    fn name(&self) -> &str {
+        "operator_to_function"
+    }
+
+    fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result<LogicalPlan> {
+        analyze_internal(plan)
+    }
+}
+
+fn analyze_internal(plan: LogicalPlan) -> Result<LogicalPlan> {
+    // OperatorToFunction is only applied to Projection
+    match plan {
+        LogicalPlan::Projection(_) => {}
+        _ => {
+            return Ok(plan);
+        }
+    }
+
+    let mut expr_rewriter = OperatorToFunctionRewriter {};
+
+    let new_expr = plan
+        .expressions()
+        .into_iter()
+        .map(|expr| expr.rewrite(&mut expr_rewriter))
+        .collect::<Result<Vec<_>>>()?;
+
+    // Not found cases that inputs more than one
+    assert_eq!(plan.inputs().len(), 1);
+    let input = plan.inputs()[0];
+
+    Ok(LogicalPlan::Projection(Projection::try_new(
+        new_expr,
+        input.to_owned().into(),
+    )?))
+}
+
+pub(crate) struct OperatorToFunctionRewriter {}
+
+impl TreeNodeRewriter for OperatorToFunctionRewriter {
+    type N = Expr;
+
+    fn mutate(&mut self, expr: Expr) -> Result<Expr> {
+        match expr {
+            Expr::BinaryExpr(BinaryExpr {
+                ref left,
+                op,
+                ref right,
+            }) => {
+                if let Some(fun) = rewrite_array_concat_operator_to_func(
+                    left.as_ref(),
+                    op,
+                    right.as_ref(),
+                ) {
+                    // Convert &Box<Expr> -> Expr
+                    let left = (**left).clone();
+                    let right = (**right).clone();
+                    return Ok(Expr::ScalarFunction(ScalarFunction {
+                        func_def: ScalarFunctionDefinition::BuiltIn(fun),
+                        args: vec![left, right],
+                    }));
+                }
+                Ok(expr)
+            }
+            _ => Ok(expr),
+        }
+    }
+}
+
+/// Summary of the logic below:
+///
+/// array || array -> array concat
+///
+/// array || scalar -> array append
+///
+/// scalar || array -> array prepend
+///
+/// (arry concat, array append, array prepend) || array -> array concat
+///
+/// (arry concat, array append, array prepend) || scalar -> array append
+fn rewrite_array_concat_operator_to_func(
+    left: &Expr,
+    op: Operator,
+    right: &Expr,
+) -> Option<BuiltinScalarFunction> {
+    // Convert `Array StringConcat Array` to ScalarFunction::ArrayConcat
+
+    if op != Operator::StringConcat {
+        return None;
+    }
+
+    match (left, right) {
+        // Chain concat operator (a || b) || array,
+        // (arry concat, array append, array prepend) || array -> array concat
+        (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::ArrayConcat),
+                args: _left_args,
+            }),
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _right_args,
+            }),
+        )
+        | (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::ArrayAppend),
+                args: _left_args,
+            }),
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _right_args,
+            }),
+        )
+        | (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::ArrayPrepend),
+                args: _left_args,
+            }),
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _right_args,
+            }),
+        ) => Some(BuiltinScalarFunction::ArrayConcat),
+        // Chain concat operator (a || b) || scalar,
+        // (arry concat, array append, array prepend) || scalar -> array append
+        (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::ArrayConcat),
+                args: _left_args,
+            }),
+            _scalar,
+        )
+        | (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::ArrayAppend),
+                args: _left_args,
+            }),
+            _scalar,
+        )
+        | (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::ArrayPrepend),
+                args: _left_args,
+            }),
+            _scalar,
+        ) => Some(BuiltinScalarFunction::ArrayAppend),
+        // array || array -> array concat
+        (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _left_args,
+            }),
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _right_args,
+            }),
+        ) => Some(BuiltinScalarFunction::ArrayConcat),
+        // array || scalar -> array append
+        (
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _left_args,
+            }),
+            _right_scalar,
+        ) => Some(BuiltinScalarFunction::ArrayAppend),
+        // scalar || array -> array prepend
+        (
+            _left_scalar,
+            Expr::ScalarFunction(ScalarFunction {
+                func_def:
+                    ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::MakeArray),
+                args: _right_args,
+            }),
+        ) => Some(BuiltinScalarFunction::ArrayPrepend),
+
+        _ => None,
+    }
+}

datafusion/physical-expr/src/expressions/binary.rs

@@ -20,9 +20,7 @@ mod kernels;
 use std::hash::{Hash, Hasher};
 use std::{any::Any, sync::Arc};
 
-use crate::array_expressions::{
-    array_append, array_concat, array_has_all, array_prepend,
-};
+use crate::array_expressions::array_has_all;
 use crate::expressions::datum::{apply, apply_cmp};
 use crate::intervals::cp_solver::{propagate_arithmetic, propagate_comparison};
 use crate::physical_expr::down_cast_any_ref;
@@ -598,12 +596,7 @@ impl BinaryExpr {
             BitwiseXor => bitwise_xor_dyn(left, right),
             BitwiseShiftRight => bitwise_shift_right_dyn(left, right),
             BitwiseShiftLeft => bitwise_shift_left_dyn(left, right),
-            StringConcat => match (left_data_type, right_data_type) {
-                (DataType::List(_), DataType::List(_)) => array_concat(&[left, right]),
-                (DataType::List(_), _) => array_append(&[left, right]),
-                (_, DataType::List(_)) => array_prepend(&[left, right]),
-                _ => binary_string_array_op!(left, right, concat_elements),
-            },
+            StringConcat => binary_string_array_op!(left, right, concat_elements),
             AtArrow => array_has_all(&[left, right]),
             ArrowAt => array_has_all(&[right, left]),
         }

datafusion/sql/src/expr/mod.rs

@@ -98,11 +98,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 StackEntry::Operator(op) => {
                     let right = eval_stack.pop().unwrap();
                     let left = eval_stack.pop().unwrap();
+
                     let expr = Expr::BinaryExpr(BinaryExpr::new(
                         Box::new(left),
                         op,
                         Box::new(right),
                     ));
+
                     eval_stack.push(expr);
                 }
             }

datafusion/sqllogictest/test_files/array.slt

@@ -4191,6 +4191,28 @@ select 1 || make_array(2, 3, 4), 1.0 || make_array(2.0, 3.0, 4.0), 'h' || make_a
 ----
 [1, 2, 3, 4] [1.0, 2.0, 3.0, 4.0] [h, e, l, l, o]
 
+# array concatenate operator with scalars #4 (mixed)
+query ?
+select 0 || [1,2,3] || 4 || [5] || [6,7];
+----
+[0, 1, 2, 3, 4, 5, 6, 7]
+
+# array concatenate operator with nd-list #5 (mixed)
+query ?
+select 0 || [1,2,3] || [[4,5]] || [[6,7,8]] || [9,10];
+----
+[[0, 1, 2, 3], [4, 5], [6, 7, 8], [9, 10]]
+
+# array concatenate operator non-valid cases
+## concat 2D with scalar is not valid
+query error
+select 0 || [1,2,3] || [[4,5]] || [[6,7,8]] || [9,10] || 11;
+
+## concat scalar with 2D is not valid
+query error
+select 0 || [[1,2,3]];
+
+
 ## array containment operator
 
 # array containment operator with scalars #1 (at arrow)