(datafusion-comet) branch main updated: fix: fall back for concat_ws with array arguments instead of failing natively (#5679)

[su...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git


The following commit(s) were added to refs/heads/main by this push:
     new 296470cab6 fix: fall back for concat_ws with array arguments instead of failing natively (#5679)
296470cab6 is described below

commit 296470cab6c3da3aec9aaf5474460b14f8a9f710
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Sat Sep 5 00:52:57 2026 +0800

    fix: fall back for concat_ws with array arguments instead of failing natively (#5679)
    
    Spark's ConcatWs accepts array<string> arguments after the separator and
    flattens their elements into the strings to join (skipping null elements).
    CometConcatWs never inspected child types and lowered every non-foldable
    call to DataFusion's concat_ws, which accepts only string arguments, so
    `concat_ws(',', arr, s)` failed at native execution with "Input was
    List(...) which is not a supported datatype for concat_ws function".
    
    CometConcatWs now mixes in CodegenDispatchFallback and returns Unsupported
    whenever any argument has an ArrayType, so these calls run through the JVM
    codegen dispatcher (Spark's own ConcatWs.doGenCode inside the Comet
    pipeline) and fall back to Spark only when the dispatcher is disabled.
    Plain string arguments keep using the native concat_ws path. A native
    kernel that flattens list arguments like Spark is left as a follow-up.
    
    Closes #5675
    
    Co-authored-by: Claude Fable 5.1 <no...@anthropic.com>
---
 docs/source/user-guide/latest/expressions.md       |  2 +-
 .../scala/org/apache/comet/serde/strings.scala     | 24 +++++++++--
 .../sql-tests/expressions/string/concat_ws.sql     | 44 +++++++++++++++++++-
 .../apache/comet/CometStringExpressionSuite.scala  | 47 +++++++++++++++++++++-
 4 files changed, 111 insertions(+), 6 deletions(-)

diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 9a59f040bc..f9818cbd73 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -561,7 +561,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
 | `chr` | ✅ | Native |  |
 | `collate` | 🔜 | — | Spark collation (umbrella [#2190](https://github.com/apache/datafusion-comet/issues/2190)) |
 | `collation` | ✅ | — | Constant-folded to a literal (Spark 4.0+) |
-| `concat_ws` | ✅ | Native |  |
+| `concat_ws` | ✅ | Hybrid | Array arguments route through the JVM codegen dispatcher; string arguments run natively ([details](compatibility/expressions/string.md)) |
 | `contains` | ✅ | — |  |
 | `decode` | ✅ | — |  |
 | `elt` | ✅ | Codegen dispatch |  |
diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala
index 665c0d3b54..278a22f1a3 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -20,7 +20,7 @@
 package org.apache.comet.serde
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, BitLength, Cast, Concat, ConcatWs, Contains, Elt, Empty2Null, EndsWith, Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, InitCap, Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay, RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike, SoundEx, StartsWith, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTranslate, Substr [...]
-import org.apache.spark.sql.types.{BinaryType, DataTypes, IntegerType, LongType, StringType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, DataTypes, IntegerType, LongType, StringType}
 
 import org.apache.comet.CometConf
 import org.apache.comet.serde.ExprOuterClass.Expr
@@ -286,12 +286,30 @@ object CometConcat
   }
 }
 
-object CometConcatWs extends CometExpressionSerde[ConcatWs] {
+object CometConcatWs extends CometExpressionSerde[ConcatWs] with CodegenDispatchFallback {
+
+  // Spark's ConcatWs accepts `array<string>` arguments after the separator and flattens their
+  // elements into the list of strings to join (null elements are skipped, like null strings).
+  // DataFusion's `concat_ws` accepts only string arguments and fails at execution time when it is
+  // handed a list, so these calls have no native path. Because this serde mixes in
+  // `CodegenDispatchFallback`, they run through the JVM codegen dispatcher (Spark's own
+  // `ConcatWs.doGenCode` inside the Comet pipeline) instead, and fall back to Spark only when the
+  // dispatcher is disabled. See https://github.com/apache/datafusion-comet/issues/5675.
+  private val arrayArgumentReason =
+    "`concat_ws` with `array<string>` arguments: Spark flattens the array elements into the " +
+      "strings to join, which DataFusion's `concat_ws` does not support " +
+      "(https://github.com/apache/datafusion-comet/issues/5675)"
+
+  override def getUnsupportedReasons(): Seq[String] = Seq(arrayArgumentReason)
 
   override def getSupportLevel(expr: ConcatWs): SupportLevel = expr.children.headOption match {
     // A NULL separator converts directly to a NULL result, so it stays supported.
     case Some(Literal(null, _)) => Compatible()
-    // Fall back to Spark for all-literal args so ConstantFolding can handle it.
+    case _ if expr.children.exists(_.dataType.isInstanceOf[ArrayType]) =>
+      Unsupported(Some(arrayArgumentReason))
+    // Decline all-literal args so that Spark's ConstantFolding handles them (it normally folds
+    // them before they reach Comet). With the dispatcher enabled they run through Spark's own
+    // generated code in-pipeline; otherwise the projection falls back to Spark.
     case _ if expr.children.forall(_.foldable) =>
       Unsupported(Some("all arguments are foldable"))
     case _ => Compatible()
diff --git a/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql b/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql
index 3de4383e86..fe6df743f7 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ws.sql
@@ -40,6 +40,48 @@ INSERT INTO names VALUES(1, 'James', 'B', 'Taylor'), (2, 'Smith', 'C', 'Davis'),
 query
 SELECT concat_ws(' ', first_name, middle_initial, last_name) FROM names
 
--- literal + literal + literal (falls back to Spark when all args are foldable)
+-- literal + literal + literal (declined natively when all args are foldable; routed through the
+-- JVM codegen dispatcher, or falls back to Spark when the dispatcher is disabled)
 query spark_answer_only
 SELECT concat_ws(',', 'hello', 'world'), concat_ws(',', '', ''), concat_ws(',', NULL, 'b', 'c'), concat_ws(NULL, 'a', 'b')
+
+-- https://github.com/apache/datafusion-comet/issues/5675
+-- Spark accepts array<string> arguments after the separator and flattens their elements into the
+-- strings to join (skipping null elements). DataFusion's concat_ws rejects list arguments, so
+-- ConcatWs mixes in CodegenDispatchFallback: these calls route through the JVM codegen dispatcher
+-- (Spark's own ConcatWs.doGenCode inside the Comet pipeline) and stay native while matching Spark
+-- exactly.
+statement
+CREATE TABLE test_concat_ws_array(arr array<string>, s string) USING parquet
+
+statement
+INSERT INTO test_concat_ws_array VALUES (array('a', 'b'), 'c d'), (array('x', NULL, 'y'), 'z'), (array('only'), ''), (CAST(array() AS array<string>), 'w'), (NULL, 'v'), (array('p', 'q'), NULL)
+
+query
+SELECT concat_ws(',', arr, s) FROM test_concat_ws_array
+
+query
+SELECT concat_ws(',', s, arr) FROM test_concat_ws_array
+
+-- single array argument
+query
+SELECT concat_ws(',', arr) FROM test_concat_ws_array
+
+-- the same array argument twice
+query
+SELECT concat_ws('-', arr, s, arr) FROM test_concat_ws_array
+
+-- array produced by another expression
+query
+SELECT concat_ws(',', split(s, ' ')) FROM test_concat_ws_array
+
+query
+SELECT concat_ws(',', split(s, ' '), arr) FROM test_concat_ws_array
+
+-- a NULL separator is NULL regardless of the argument types and stays on the native path
+query
+SELECT concat_ws(NULL, arr, s) FROM test_concat_ws_array
+
+-- literal array arguments
+query
+SELECT concat_ws(',', array('a', 'b'), 'c'), concat_ws(',', array('x', NULL, 'y')), concat_ws(',', CAST(array() AS array<string>), 'w')
diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
index a0fa91b8fe..bbe941417b 100644
--- a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
@@ -27,8 +27,9 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
 
 import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
+import org.apache.comet.udf.codegen.CometScalaUDFCodegen
 
-class CometStringExpressionSuite extends CometTestBase {
+class CometStringExpressionSuite extends CometTestBase with CometCodegenAssertions {
   // scalastyle:off
   private val edgeCases = Seq(
     "é", // unicode 'e\\u{301}'
@@ -733,4 +734,48 @@ class CometStringExpressionSuite extends CometTestBase {
     // scalastyle:on
   }
 
+  test("concat_ws with array<string> arguments") {
+    // https://github.com/apache/datafusion-comet/issues/5675
+    // Spark flattens array<string> arguments into the strings to join (skipping null elements).
+    // DataFusion's concat_ws rejects list arguments, so these calls run through the JVM codegen
+    // dispatcher (Spark's own doGenCode inside the Comet pipeline) instead of the native path.
+    val data: Seq[(Seq[String], String)] = Seq(
+      (Seq("a", "b"), "c d"),
+      (Seq("x", null, "y"), "z"),
+      (Seq("only"), ""),
+      (Seq.empty[String], "w"),
+      (null, "v"),
+      (Seq("p", "q"), null))
+    withParquetTable(data, "tbl") {
+      val arrayArgQueries = Seq(
+        "SELECT concat_ws(',', _1, _2) FROM tbl",
+        "SELECT concat_ws(',', _2, _1) FROM tbl",
+        "SELECT concat_ws(',', _1) FROM tbl",
+        "SELECT concat_ws('-', _1, _2, _1) FROM tbl",
+        "SELECT concat_ws(',', split(_2, ' ')) FROM tbl",
+        "SELECT concat_ws(',', split(_2, ' '), _1) FROM tbl")
+      for (query <- arrayArgQueries) {
+        // Spark's answer, the whole plan stays in Comet, and the codegen dispatcher actually ran.
+        assertCodegenRan {
+          checkSparkAnswerAndOperator(query)
+        }
+      }
+      // With the dispatcher disabled there is no in-pipeline path, so the projection falls back
+      // to Spark with the serde's reason instead of failing at native execution.
+      withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+        for (query <- arrayArgQueries) {
+          checkSparkAnswerAndFallbackReason(query, "`concat_ws` with `array<string>` arguments")
+        }
+      }
+      // A NULL separator produces NULL regardless of the argument types and stays native.
+      checkSparkAnswerAndOperator("SELECT concat_ws(NULL, _1, _2) FROM tbl")
+      // Plain string arguments keep using the native concat_ws path, not the dispatcher.
+      CometScalaUDFCodegen.resetStats()
+      checkSparkAnswerAndOperator("SELECT concat_ws(',', _2, 'lit', _2) FROM tbl")
+      assert(
+        CometScalaUDFCodegen.stats().totalLookups == 0,
+        "expected the native concat_ws path for string arguments, not codegen dispatch")
+    }
+  }
+
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org