From 038b498456bf94bbd5103237d6baf07dd607829e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 4 Sep 2026 20:29:03 +0800 Subject: [PATCH] fix: fall back for concat_ws with array arguments instead of failing natively Spark's ConcatWs accepts array 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 --- docs/source/user-guide/latest/expressions.md | 2 +- .../org/apache/comet/serde/strings.scala | 24 ++++++++-- .../expressions/string/concat_ws.sql | 44 ++++++++++++++++- .../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 12c7f47f069..fe393e8f413 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 665c0d3b54e..278a22f1a3f 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, Substring, SubstringIndex, ToCharacter, ToNumber, TryToNumber, UnBase64, Upper} -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` 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` 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 3de4383e863..fe6df743f7a 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 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, 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), '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), 'w') diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala index 0b5acc70a28..45e5133389f 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( "eĢ", // unicode 'e\\u{301}' @@ -712,4 +713,48 @@ class CometStringExpressionSuite extends CometTestBase { // scalastyle:on } + test("concat_ws with array arguments") { + // https://github.com/apache/datafusion-comet/issues/5675 + // Spark flattens array 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` 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") + } + } + }