feat: default native replace for non-empty UTF8_BINARY literal search

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 `src` and `replace` are each a short well-formed literal or a column, and `search` is a short, well-formed, non-empty `UTF8_BINARY` literal.
+  - Nested source/replacement expressions (`substring`, `concat`, casts that can throw), empty / malformed / oversized literals, 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

native/core/src/execution/operators/batch_split.rs

@@ -0,0 +1,154 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+use datafusion::common::{internal_err, Result};
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet, SplitMetrics};
+use datafusion::physical_plan::stream::BatchSplitStream;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
+};
+
+/// Splits oversized batches emitted by an input plan using the runtime batch size.
+///
+/// DataFusion 54.1.0's `UnnestExec` can emit more rows than the configured batch size. This
+/// wrapper is a temporary downstream boundary for Comet's explode path until Comet upgrades to a
+/// DataFusion version containing https://github.com/apache/datafusion/pull/24384.
+#[derive(Debug)]
+pub struct BatchSplitExec {
+    input: Arc<dyn ExecutionPlan>,
+    cache: Arc<PlanProperties>,
+    metrics: ExecutionPlanMetricsSet,
+}
+
+impl BatchSplitExec {
+    pub fn new(input: Arc<dyn ExecutionPlan>) -> Self {
+        Self {
+            cache: Arc::clone(input.properties()),
+            input,
+            metrics: ExecutionPlanMetricsSet::new(),
+        }
+    }
+}
+
+impl DisplayAs for BatchSplitExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "CometBatchSplitExec")
+            }
+            DisplayFormatType::TreeRender => unimplemented!(),
+        }
+    }
+}
+
+impl ExecutionPlan for BatchSplitExec {
+    fn name(&self) -> &str {
+        "CometBatchSplitExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.cache
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return internal_err!(
+                "CometBatchSplitExec expects one child, got {}",
+                children.len()
+            );
+        }
+        Ok(Arc::new(Self::new(Arc::clone(&children[0]))))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let batch_size = context.session_config().batch_size();
+        let input = self.input.execute(partition, context)?;
+        Ok(Box::pin(BatchSplitStream::new(
+            input,
+            batch_size,
+            SplitMetrics::new(&self.metrics, partition),
+        )))
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        vec![true]
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use arrow::array::{AsArray, Int32Array, RecordBatch};
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::datasource::memory::MemorySourceConfig;
+    use datafusion::datasource::source::DataSourceExec;
+    use datafusion::execution::SessionStateBuilder;
+    use datafusion::prelude::{SessionConfig, SessionContext};
+    use futures::StreamExt;
+
+    #[tokio::test]
+    async fn splits_batches_at_the_runtime_batch_size_and_preserves_order() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let input_batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int32Array::from_iter_values(0..10))],
+        )
+        .unwrap();
+        let source = MemorySourceConfig::try_new(&[vec![input_batch]], schema, None).unwrap();
+        let input: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(source)));
+        let split = BatchSplitExec::new(input);
+
+        let config = SessionConfig::new().with_batch_size(4);
+        let state = SessionStateBuilder::new().with_config(config).build();
+        let context = SessionContext::new_with_state(state);
+        let mut stream = split.execute(0, context.task_ctx()).unwrap();
+
+        let mut batch_sizes = vec![];
+        let mut values = vec![];
+        while let Some(batch) = stream.next().await {
+            let batch = batch.unwrap();
+            batch_sizes.push(batch.num_rows());
+            let column = batch
+                .column(0)
+                .as_primitive::<arrow::datatypes::Int32Type>();
+            values.extend(column.values().iter().copied());
+        }
+
+        assert_eq!(batch_sizes, vec![4, 4, 2]);
+        assert_eq!(values, (0..10).collect::<Vec<_>>());
+    }
+}

native/core/src/execution/operators/mod.rs

@@ -20,11 +20,13 @@
 pub use crate::errors::ExecutionError;
 
 pub use aligned_stream_reader::*;
+pub use batch_split::BatchSplitExec;
 pub use copy::*;
 pub use iceberg_scan::*;
 pub use scan::*;
 
 mod aligned_stream_reader;
+mod batch_split;
 mod copy;
 mod expand;
 pub use expand::ExpandExec;

native/core/src/execution/planner.rs

@@ -38,8 +38,8 @@ use crate::execution::{
     expressions::list_positions::ListPositionsExpr,
     expressions::subquery::Subquery,
     operators::{
-        ExecutionError, ExpandExec, ParquetCompression, ParquetWriterExec, SampleExec, ScanExec,
-        ShuffleScanExec,
+        BatchSplitExec, ExecutionError, ExpandExec, ParquetCompression, ParquetWriterExec,
+        SampleExec, ScanExec, ShuffleScanExec,
     },
     planner::expression_registry::ExpressionRegistry,
     planner::operator_registry::OperatorRegistry,
@@ -2116,11 +2116,20 @@ impl PhysicalPlanner {
                     output_schema,
                     unnest_options,
                 )?);
+                // DataFusion 54.1.0's UnnestExec can emit more than the runtime batch size.
+                // Bound batches before downstream native projections until Comet upgrades to a
+                // DataFusion version containing https://github.com/apache/datafusion/pull/24384.
+                let bounded_unnest: Arc<dyn ExecutionPlan> =
+                    Arc::new(BatchSplitExec::new(unnest_exec));
 
                 Ok((
                     scans,
                     shuffle_scans,
-                    Arc::new(SparkPlan::new(spark_plan.plan_id, unnest_exec, vec![child])),
+                    Arc::new(SparkPlan::new(
+                        spark_plan.plan_id,
+                        bounded_unnest,
+                        vec![child],
+                    )),
                 ))
             }
             OpStruct::SortMergeJoin(join) => {

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,93 @@ 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. `src` uses the same whitelist as `replace`: a short well-formed literal or a
+   * column. Nested expressions (`substring`, `concat`, …) stay on the dispatcher so a throwing or
+   * malformed child cannot hide inside the source tree. 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 sourceIsSafe = isNativeSafeStringArg(children(0), allowEmptyLiteral = true)
+    val searchIsSafe = children(1) match {
+      case Literal(v: UTF8String, _) => isNativeSafeStringLiteral(v, allowEmpty = false)
+      case _ => false
+    }
+    val replacementIsSafe = isNativeSafeStringArg(children(2), allowEmptyLiteral = true)
+    val utf8BinaryCollation =
+      !children.exists(c => QueryPlanSerde.isStringCollationType(c.dataType))
+    utf8BinaryCollation && sourceIsSafe && searchIsSafe && replacementIsSafe
+  }
+
+  /** A column, null literal, or a short well-formed string literal. */
+  private def isNativeSafeStringArg(expr: Expression, allowEmptyLiteral: Boolean): Boolean =
+    expr match {
+      case Literal(null, _) => true
+      case Literal(v: UTF8String, _) => isNativeSafeStringLiteral(v, allowEmptyLiteral)
+      case _: Attribute | _: BoundReference => true
+      case _ => false
+    }
+
+  /**
+   * `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. Native operators feeding projections must emit batches no
+   * larger than `spark.comet.batchSize`; Comet's explode path enforces that with
+   * `BatchSplitExec`.
+   */
+  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 `src` and `replace` are each a short well-formed literal or a column, and " +
+        "`search` is a short, well-formed, non-empty `UTF8_BINARY` literal, 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.
+      // Nested / malformed / oversized / throwing children, 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/resources/sql-tests/expressions/string/string_replace.sql

@@ -19,15 +19,16 @@ 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
 
 -- 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