fix: narrow replace native-safe subset for expression-boundary cases

Author: Hung

docs/source/contributor-guide/expression-audits/string_funcs.md

@@ -169,8 +169,8 @@
 
 - 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.
-  - Comet evaluates `replace` natively by default when `search` is a non-empty `UTF8_BINARY` literal.
-  - Empty literal search, non-literal search, and non-default collations stay on the JVM codegen dispatcher.
+  - 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)).
 

spark/src/main/scala/org/apache/comet/serde/strings.scala

@@ -19,7 +19,10 @@
 
 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
 
@@ -180,37 +183,66 @@ object CometStringReplace
     with NativeOptInAvailable {
 
   /**
-   * Native DataFusion `replace` differs from Spark only when the search string is empty (Spark
-   * returns `src` unchanged; DataFusion inserts the replacement between every character). That
-   * case is decidable at plan time when `search` is a literal.
+   * 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 native kernel is also byte-level `UTF8_BINARY` only, so non-default collations stay on
-   * the dispatcher. https://github.com/apache/datafusion-comet/issues/4496
+   * 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 nativeSafeSearchSubset(expr: StringReplace): Boolean = {
+  private def nativeSafeSubset(expr: StringReplace): Boolean = {
     val children = expr.children
     if (children.length != 3) {
       return false
     }
-    val searchIsNonEmptyLiteral = children(1) match {
-      case Literal(v: UTF8String, _) => v != null && v.numBytes() > 0
+    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 && searchIsNonEmptyLiteral
+    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 non-empty `UTF8_BINARY` literal, Comet evaluates `replace` natively " +
+      "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)) || nativeSafeSearchSubset(expr)) {
+    if (CometConf.isExprAllowIncompat(getExprConfigName(expr)) || nativeSafeSubset(expr)) {
       Compatible()
     } else {
       Compatible(nativeOptIn =
@@ -221,14 +253,13 @@ object CometStringReplace
       expr: StringReplace,
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] = {
-    if (CometConf.isExprAllowIncompat(getExprConfigName(expr)) || nativeSafeSearchSubset(expr)) {
-      // Native DataFusion `replace` matches Spark when search is a non-empty UTF8_BINARY
-      // literal (the common case, selected by default) and when the user has opted in.
+    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 {
-      // Empty literal search, non-literal search, 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.
+      // 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)
     }
   }

spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala

@@ -331,8 +331,8 @@ 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) USING parquet")
-      sql("INSERT INTO t VALUES ('hello world', 'world'), ('abcabc', 'world')")
+      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",
@@ -349,6 +349,14 @@ class CometCodegenSuite
           !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 =
@@ -381,6 +389,62 @@ class CometCodegenSuite
     }
   }
 
+  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 `<name>: <reason>` so distinct expressions
     // don't collapse in the `Set[String]` roll-up.