feat: support concat for strings

Author: comphead

docs/source/user-guide/latest/configs.md

@@ -200,6 +200,7 @@ These settings can be used to determine which parts of the plan are accelerated
 | `spark.comet.expression.CheckOverflow.enabled` | Enable Comet acceleration for `CheckOverflow` | true |
 | `spark.comet.expression.Chr.enabled` | Enable Comet acceleration for `Chr` | true |
 | `spark.comet.expression.Coalesce.enabled` | Enable Comet acceleration for `Coalesce` | true |
+| `spark.comet.expression.Concat.enabled` | Enable Comet acceleration for `Concat` | true |
 | `spark.comet.expression.ConcatWs.enabled` | Enable Comet acceleration for `ConcatWs` | true |
 | `spark.comet.expression.Contains.enabled` | Enable Comet acceleration for `Contains` | true |
 | `spark.comet.expression.Cos.enabled` | Enable Comet acceleration for `Cos` | true |

native/Cargo.lock

@@ -47,6 +47,7 @@ use datafusion_spark::function::hash::sha1::SparkSha1;
 use datafusion_spark::function::hash::sha2::SparkSha2;
 use datafusion_spark::function::math::expm1::SparkExpm1;
 use datafusion_spark::function::string::char::CharFunc;
+use datafusion_spark::function::string::concat::SparkConcat;
 use futures::poll;
 use futures::stream::StreamExt;
 use jni::objects::JByteBuffer;
@@ -317,20 +318,23 @@ fn prepare_datafusion_session_context(
     let mut session_ctx = SessionContext::new_with_config_rt(session_config, Arc::new(runtime));
 
     datafusion::functions_nested::register_all(&mut session_ctx)?;
+    register_datafusion_spark_function(&session_ctx);
+    // Must be the last one to override existing functions with the same name
+    datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?;
+
+    Ok(session_ctx)
+}
 
-    // register UDFs from datafusion-spark crate
+// register UDFs from datafusion-spark crate
+fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkExpm1::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha2::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(CharFunc::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitGet::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateAdd::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateSub::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha1::default()));
-
-    // Must be the last one to override existing functions with the same name
-    datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?;
-
-    Ok(session_ctx)
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default()));
 }
 
 /// Prepares arrow arrays for output.

native/core/src/execution/jni_api.rs

@@ -30,6 +30,7 @@ edition = { workspace = true }
 arrow = { workspace = true }
 chrono = { workspace = true }
 datafusion = { workspace = true }
+datafusion-spark = { workspace = true }
 chrono-tz = { workspace = true }
 num = { workspace = true }
 regex = { workspace = true }

native/spark-expr/Cargo.toml

@@ -180,6 +180,17 @@ pub fn create_comet_physical_fun_with_eval_mode(
             let func = Arc::new(spark_modulo);
             make_comet_scalar_udf!("spark_modulo", func, without data_type, fail_on_error)
         }
+        "concat" => {
+            // Use concat from datafusion-spark crate which has Spark semantics
+            // (returns null if any argument is null)
+            // The registry parameter already contains functions from datafusion-spark
+            // if they were registered in prepare_datafusion_session_context
+            registry.udf("concat").map_err(|e| {
+                DataFusionError::Execution(format!(
+                    "Function concat not found in the registry: {e}",
+                ))
+            })
+        }
         _ => registry.udf(fun_name).map_err(|e| {
             DataFusionError::Execution(format!(
                 "Function {fun_name} not found in the registry: {e}",

native/spark-expr/src/comet_scalar_funcs.rs

@@ -3149,10 +3149,39 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       sql(
         "create table t1 using parquet as select uuid() c1, uuid() c2, uuid() c3, uuid() c4, cast(null as string) c5 from range(10)")
       checkSparkAnswerAndOperator("select concat(c1, c2) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c1) AS x FROM t1")
       checkSparkAnswerAndOperator("select concat(c1, c2, c3) AS x FROM t1")
       checkSparkAnswerAndOperator("select concat(c1, c2, c3, c5) AS x FROM t1")
-      //checkSparkAnswerAndOperator("select concat(concat(c1, c2, c3), concat(c1, c3)) AS x FROM t1")
+      checkSparkAnswerAndOperator(
+        "select concat(concat(c1, c2, c3), concat(c1, c3)) AS x FROM t1")
+    }
+  }
+
+  // https://github.com/apache/datafusion-comet/issues/2647
+  ignore("test concat function - arrays") {
+    withTable("t1") {
+      sql(
+        "create table t1 using parquet as select array(id, id+1) c1, array(id+2, id+3) c2, array() c3, array(null) c4, cast(null as array<int>) c5 from range(10)")
+      checkSparkAnswerAndOperator("select concat(c1, c2) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c1) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c2, c3) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c2, c3, c5) AS x FROM t1")
+      checkSparkAnswerAndOperator(
+        "select concat(concat(c1, c2, c3), concat(c1, c3)) AS x FROM t1")
     }
   }
 
+  // https://github.com/apache/datafusion-comet/issues/2647
+  ignore("test concat function - binary") {
+    withTable("t1") {
+      sql(
+        "create table t1 using parquet as select cast(uuid() as binary) c1, cast(uuid() as binary) c2, cast(uuid() as binary) c3, cast(uuid() as binary) c4, cast(null as binary) c5 from range(10)")
+      checkSparkAnswerAndOperator("select concat(c1, c2) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c1) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c2, c3) AS x FROM t1")
+      checkSparkAnswerAndOperator("select concat(c1, c2, c3, c5) AS x FROM t1")
+      checkSparkAnswerAndOperator(
+        "select concat(concat(c1, c2, c3), concat(c1, c3)) AS x FROM t1")
+    }
+  }
 }