diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 12c7f47f06..fe393e8f41 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, 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 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 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 0b5acc70a2..45e5133389 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") + } + } + }