diff --git a/docs/source/contributor-guide/expression-audits/string_funcs.md b/docs/source/contributor-guide/expression-audits/string_funcs.md index 58b300ebc12..2159cdc66b0 100644 --- a/docs/source/contributor-guide/expression-audits/string_funcs.md +++ b/docs/source/contributor-guide/expression-audits/string_funcs.md @@ -168,7 +168,10 @@ ## replace - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `StringReplace(src, search, replace)`; when `search` is empty, Spark returns `src` unchanged (short-circuit on `search.numBytes == 0`). DataFusion `replace` instead inserts `replace` between every character, so `CometStringReplace` reports `Compatible` with a `NativeOptIn` and runs Spark's own generated code inside the Comet pipeline by default. The native DataFusion `replace` is used only when `spark.comet.expression.StringReplace.allowIncompatible=true`. +- Spark 3.5.8 (audited 2026-05-27): baseline. `StringReplace(src, search, replace)`; when `search` is empty, Spark returns `src` unchanged (short-circuit on `search.numBytes == 0`). DataFusion `replace` instead inserts `replace` between every character. + - Comet evaluates `replace` natively by default when `search` is a short, well-formed, non-empty `UTF8_BINARY` literal and `replace` is a short well-formed literal or a column. + - Empty, malformed, or oversized literals, a non-literal / non-column replacement (including expressions that can throw), and non-default collations stay on the JVM codegen dispatcher. The kernel match is not enough: `CometLiteral` is not byte-preserving, DataFusion evaluates every child before `replace`, and large scalars overflow Arrow `Utf8` offsets when broadcast. + - Users can still opt into the native (potentially incompatible) path for remaining cases via `spark.comet.expression.StringReplace.allowIncompatible=true`. - Spark 4.0.1 (audited 2026-05-27): routes through `CollationSupport.StringReplace.exec`; semantics unchanged for `UTF8_BINARY`. Non-default collations not honoured by Comet ([#4496](https://github.com/apache/datafusion-comet/issues/4496)). ## right 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 ebf45089882..618711136f3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -19,8 +19,12 @@ 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 java.nio.charset.StandardCharsets +import java.util.Arrays + +import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, BitLength, BoundReference, 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.unsafe.types.UTF8String import org.apache.comet.CometConf import org.apache.comet.serde.ExprOuterClass.Expr @@ -178,29 +182,84 @@ object CometStringReplace extends CometScalarFunction[StringReplace]("replace") with NativeOptInAvailable { + /** + * The DataFusion `replace` kernel matches Spark only for a non-empty search string. Kernel + * compatibility is not enough: `CometLiteral` serializes strings via `UTF8String.toString` + * (malformed UTF-8 becomes U+FFFD), DataFusion evaluates every child before `replace` (so a + * NULL `src` does not skip a throwing replacement), and scalar literals are broadcast into + * Arrow `Utf8` arrays that overflow 32-bit offsets on a large batch. + * + * The default native path is therefore limited to a plan-time subset that avoids those + * boundaries. Non-default collations stay on the dispatcher. + * https://github.com/apache/datafusion-comet/issues/4496 + */ + private def nativeSafeSubset(expr: StringReplace): Boolean = { + val children = expr.children + if (children.length != 3) { + return false + } + val searchIsSafe = children(1) match { + case Literal(v: UTF8String, _) => isNativeSafeStringLiteral(v, allowEmpty = false) + case _ => false + } + val replacementIsSafe = children(2) match { + case Literal(null, _) => true + case Literal(v: UTF8String, _) => isNativeSafeStringLiteral(v, allowEmpty = true) + case _: Attribute | _: BoundReference => true + case _ => false + } + val utf8BinaryCollation = + !children.exists(c => QueryPlanSerde.isStringCollationType(c.dataType)) + utf8BinaryCollation && searchIsSafe && replacementIsSafe + } + + /** + * `CometLiteral` encodes a string as `UTF8String.toString`, so only literals whose bytes + * survive that round-trip can be sent natively. The size cap keeps a broadcast scalar under + * Arrow `Utf8`'s 32-bit offset limit at `spark.comet.batchSize` rows. + */ + private def isNativeSafeStringLiteral(v: UTF8String, allowEmpty: Boolean): Boolean = { + if (v == null) { + return false + } + if (!allowEmpty && v.numBytes() == 0) { + return false + } + val maxBytes = Int.MaxValue / math.max(CometConf.COMET_BATCH_SIZE.get(), 1) + if (v.numBytes() > maxBytes) { + return false + } + Arrays.equals(v.getBytes, v.toString.getBytes(StandardCharsets.UTF_8)) + } + + override def getCompatibleNotes(): Seq[String] = + Seq( + "When `search` is a short, well-formed, non-empty `UTF8_BINARY` literal and `replace` " + + "is a short well-formed literal or a column, Comet evaluates `replace` natively " + + "by default.") + override def getIncompatibleReasons(): Seq[String] = Seq("Produces different results from Spark when the search string is empty") override def getSupportLevel(expr: StringReplace): SupportLevel = - if (!CometConf.isExprAllowIncompat(getExprConfigName(expr))) { + if (CometConf.isExprAllowIncompat(getExprConfigName(expr)) || nativeSafeSubset(expr)) { + Compatible() + } else { Compatible(nativeOptIn = Some(NativeOptIn(CometConf.getExprAllowIncompatConfigKey(getExprConfigName(expr))))) - } else { - Compatible() } override def convert( expr: StringReplace, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { - if (CometConf.isExprAllowIncompat(getExprConfigName(expr))) { - // The native DataFusion `replace` avoids the JVM allocations of the codegen - // dispatcher but is not Spark-compatible for an empty search string, so it is - // only used when incompatibility is explicitly allowed. + if (CometConf.isExprAllowIncompat(getExprConfigName(expr)) || nativeSafeSubset(expr)) { + // Native path for the plan-time safe subset, or when the user has opted in. super.convert(expr, inputs, binding) } else { - // Run Spark's own generated code inside the Comet pipeline so the result matches Spark - // exactly. Falls back to Spark when the codegen dispatcher is disabled. + // Empty / malformed / oversized search, a throwing or non-column replacement, or a + // non-default collation: run Spark's own generated code inside the Comet pipeline. + // Falls back to Spark when the codegen dispatcher is disabled. CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) } } diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_replace.sql b/spark/src/test/resources/sql-tests/expressions/string/string_replace.sql index ad09525b6d9..bd88bd89ed9 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_replace.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_replace.sql @@ -19,7 +19,7 @@ statement CREATE TABLE test_str_replace(s string, search string, replace string) USING parquet statement -INSERT INTO test_str_replace VALUES ('hello world', 'world', 'there'), ('aaa', 'a', 'bb'), ('hello', 'xyz', 'abc'), ('', 'a', 'b'), (NULL, 'a', 'b'), ('hello', '', 'x') +INSERT INTO test_str_replace VALUES ('hello world', 'world', 'there'), ('aaa', 'a', 'bb'), ('hello', 'xyz', 'abc'), ('', 'a', 'b'), (NULL, 'a', 'b'), ('hello', '', 'x'), ('aaaa', 'aa', 'x'), ('你好你好', '你好', 'X'), ('😀a😀', '😀', 'x') query SELECT replace(s, search, replace) FROM test_str_replace @@ -27,7 +27,8 @@ SELECT replace(s, search, replace) FROM test_str_replace -- Empty literal search: DataFusion's replace diverges from Spark -- (Spark short-circuits and returns the source unchanged). The custom -- CometStringReplace serde routes through the codegen dispatcher so --- Spark's own doGenCode handles this case. +-- Spark's own doGenCode handles this case. Non-empty UTF8_BINARY literal +-- search takes the native path by default (#5354). -- https://github.com/apache/datafusion-comet/issues/4497 query SELECT replace('hello', '', 'x') @@ -41,6 +42,18 @@ SELECT replace(NULL, '', 'x') query SELECT replace('hello', '', NULL) +-- Overlapping candidates: Spark replaces non-overlapping left-to-right +-- ('aaaa' + 'aa' -> 'xx'). Column source + literal search takes the native path. +query +SELECT replace(s, 'aa', 'x') FROM test_str_replace WHERE s = 'aaaa' + +-- Multi-byte UTF-8 values. Column source + literal search takes the native path. +query +SELECT replace(s, '你好', 'X') FROM test_str_replace WHERE s = '你好你好' + +query +SELECT replace(s, '😀', 'x') FROM test_str_replace WHERE s = '😀a😀' + -- column + literal + literal query SELECT replace(s, 'world', 'there') FROM test_str_replace diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5806cb35015..529836e9144 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec @@ -328,6 +329,122 @@ class CometCodegenSuite } } + test("replace routes native vs JVM codegen dispatcher based on non-empty literal search") { + withTable("t") { + sql("CREATE TABLE t (s STRING, search STRING, r STRING) USING parquet") + sql("INSERT INTO t VALUES ('hello world', 'world', 'comet'), ('abcabc', 'world', 'comet')") + + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + + val dfNative = sql("SELECT replace(s, 'world', 'comet') FROM t") + checkSparkAnswerAndOperator(dfNative) + val explainNative = + new ExtendedExplainInfo().generateExtendedInfo(dfNative.queryExecution.executedPlan) + assert( + !explainNative.contains("JVM codegen dispatcher: replace"), + s"expected native path for non-empty literal search, got:\n$explainNative") + + val dfColReplace = sql("SELECT replace(s, 'world', r) FROM t") + checkSparkAnswerAndOperator(dfColReplace) + val explainColReplace = + new ExtendedExplainInfo().generateExtendedInfo(dfColReplace.queryExecution.executedPlan) + assert( + !explainColReplace.contains("JVM codegen dispatcher: replace"), + s"expected native path for column replacement, got:\n$explainColReplace") + + val dfEmptySearch = sql("SELECT replace(s, '', 'comet') FROM t") + checkSparkAnswerAndOperator(dfEmptySearch) + val explainEmptySearch = + new ExtendedExplainInfo().generateExtendedInfo( + dfEmptySearch.queryExecution.executedPlan) + assert( + explainEmptySearch.contains("JVM codegen dispatcher: replace"), + s"expected dispatcher path for empty literal search, got:\n$explainEmptySearch") + + val dfNonLiteralSearch = sql("SELECT replace(s, search, 'comet') FROM t") + checkSparkAnswerAndOperator(dfNonLiteralSearch) + val explainNonLiteralSearch = + new ExtendedExplainInfo().generateExtendedInfo( + dfNonLiteralSearch.queryExecution.executedPlan) + assert( + explainNonLiteralSearch.contains("JVM codegen dispatcher: replace"), + s"expected dispatcher path for non-literal search, got:\n$explainNonLiteralSearch") + + if (isSpark40Plus) { + val dfCollated = + sql("SELECT replace(CAST(s AS STRING COLLATE UTF8_LCASE), 'world', 'comet') FROM t") + checkSparkAnswerAndOperator(dfCollated) + val explainCollated = + new ExtendedExplainInfo().generateExtendedInfo(dfCollated.queryExecution.executedPlan) + assert( + explainCollated.contains("JVM codegen dispatcher: replace"), + s"expected dispatcher path for non-UTF8_BINARY collation, got:\n$explainCollated") + } + } + } + } + + test("replace stays on dispatcher for expression-boundary incompatibilities") { + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + + def assertDispatcher(df: org.apache.spark.sql.DataFrame, clue: String): Unit = { + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert(explain.contains("JVM codegen dispatcher: replace"), s"$clue, got:\n$explain") + } + + // Malformed search: CometLiteral would normalize 0xFF to U+FFFD, incorrectly matching + // a well-formed U+FFFD in the source. + withTable("t") { + sql("CREATE TABLE t (s STRING) USING parquet") + sql("INSERT INTO t VALUES ('\uFFFD'), ('ok')") + assertDispatcher( + sql("SELECT replace(s, CAST(X'FF' AS STRING), 'x') FROM t"), + "expected dispatcher path for malformed search literal") + } + + // Malformed replacement has the same serialization hazard. + withTable("t") { + sql("CREATE TABLE t (s STRING) USING parquet") + sql("INSERT INTO t VALUES ('a'), ('b')") + assertDispatcher( + sql("SELECT replace(s, 'a', CAST(X'FF' AS STRING)) FROM t"), + "expected dispatcher path for malformed replacement literal") + } + + // Spark skips replacement evaluation when src is NULL; native evaluates every child. + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTable("t") { + sql("CREATE TABLE t (s STRING, n INT) USING parquet") + sql("INSERT INTO t VALUES (NULL, 0), ('a', 1)") + assertDispatcher( + sql("SELECT replace(s, 'a', CAST(1 / n AS STRING)) FROM t"), + "expected dispatcher path for throwing replacement expression") + } + } + + // A 256 KiB scalar replacement overflows Arrow Utf8 offsets when broadcast to 8192 rows. + withTable("t") { + sql("CREATE TABLE t (s STRING) USING parquet") + sql("INSERT INTO t VALUES ('hello')") + assertDispatcher( + sql("SELECT replace(s, 'notfound', repeat('x', 262144)) FROM t"), + "expected dispatcher path for oversized replacement literal") + } + } + } + test("codegen dispatch fallback reasons name the expression") { // Flag-off short-circuit tags the expression `: ` so distinct expressions // don't collapse in the `Set[String]` roll-up.