You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/19 23:36:53 UTC
(datafusion-comet) 01/01: feat: enable native replace for safe inputs (#6680)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6680-3355602e268819750884aa0ce0dd47002c6ea580
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 7eaadb5ba766027064b8b50de4c8ddb0b9b8c96f
Author: Yu Hsun Hung <75...@users.noreply.github.com>
AuthorDate: Mon Oct 19 23:36:30 2026 +0000
feat: enable native replace for safe inputs (#6680)
* Enable native replace for safe inputs
* Test native replace literal and broadcast boundaries
* docs: format string expression audit tables
* Fix replace routing tests for native safe inputs
---
.../expression-audits/string_funcs.md | 28 +++-
docs/source/user-guide/latest/expressions.md | 2 +-
.../scala/org/apache/comet/serde/strings.scala | 58 +++++++--
.../string/routing_string_opt_in_disabled.sql | 7 +-
.../string/routing_string_opt_in_enabled.sql | 7 +-
.../expressions/string/string_replace.sql | 5 +-
.../scala/org/apache/comet/CometCodegenSuite.scala | 144 +++++++++++++++++++++
.../benchmark/CometStringExpressionBenchmark.scala | 85 +++++++++++-
8 files changed, 315 insertions(+), 21 deletions(-)
diff --git a/docs/source/contributor-guide/expression-audits/string_funcs.md b/docs/source/contributor-guide/expression-audits/string_funcs.md
index 12721470b4..36cd79613f 100644
--- a/docs/source/contributor-guide/expression-audits/string_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/string_funcs.md
@@ -169,8 +169,34 @@
## 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.allowIncom [...]
+- 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`). Earlier DataFusion versions inserted `replace` between every character, so Comet routed the expression through Spark's generated code by default.
- 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)).
+- DataFusion 55.1.0 (implementation reviewed 2026-10-05, [#5354](https://github.com/apache/datafusion-comet/issues/5354)): an empty search now returns the source unchanged. A direct source column and non-null literal search/replacement use its scalar fast path, which does not broadcast the literals into Arrow `Utf8` arrays. With a replacement column, its general path broadcasts the search literal into an Arrow array, so Comet limits that search to 256 bytes and the configured batch's bro [...]
+- Benchmark (2026-10-05): `CometStringExpressionBenchmark` on Apple M4, macOS 26.7.1, JDK 21.0.12, Spark 4.1, release build, 65,536 rows, one run per route. The same query and Parquet inputs were used with only the default `replace` routing changed. These end-to-end figures include scan and projection; 10% match density is approximate.
+
+ | Match density | Dispatcher ns/row | Native ns/row | Dispatcher M rows/s | Native M rows/s | Native speedup |
+ | ------------- | ----------------: | ------------: | ------------------: | --------------: | -------------: |
+ | 0% | 207.5 | 149.5 | 4.8 | 6.7 | 1.39× |
+ | ~10% | 277.5 | 141.8 | 3.6 | 7.1 | 1.96× |
+ | 100% | 309.9 | 170.5 | 3.2 | 5.9 | 1.82× |
+
+- Column replacement benchmark rerun (2026-10-05): same Apple M4, Spark 4.1 release build and 65,536-row setup, one run. Spark versus native Comet measured 263.5 versus 149.9 ns/row (3.8 versus 6.7 M rows/s) at 0% matches, 225.2 versus 136.9 (4.4 versus 7.3 M rows/s) at about 10%, and 286.5 versus 164.9 (3.5 versus 6.1 M rows/s) at 100%. End-to-end scan and projection times.
+- Column replacement dispatcher baseline (2026-10-05): same machine and setup, with only the default route temporarily forced through the JVM codegen dispatcher. Dispatcher versus native measured 184.4 versus 149.9 ns/row (5.4 versus 6.7 M rows/s, 1.23× native speedup) at 0% matches, 243.3 versus 136.9 (4.1 versus 7.3 M rows/s, 1.78×) at about 10%, and 281.5 versus 164.9 (3.6 versus 6.1 M rows/s, 1.71×) at 100%. The dispatcher and native runs were separate single runs; benchmark timing i [...]
+- Long search literal benchmark (2026-10-05): Apple M4, macOS 26.7.1, JDK 21.0.12, Spark 4.1 release build, 65,536 rows, one run per route and case. Source strings were 2,048 bytes; search literals were 16, 256, or 1,024 ASCII bytes with 0%, about 10%, or 100% match density. Each run includes scan and projection. Dispatcher / native ns per row and native speedup over dispatcher:
+
+ | Search bytes | Match density | Dispatcher ns/row | Native ns/row | Dispatcher M rows/s | Native M rows/s | Native speedup |
+ | -----------: | ------------- | ----------------: | ------------: | ------------------: | --------------: | -------------: |
+ | 16 | 0% | 1,734.6 | 723.6 | 0.6 | 1.4 | 2.40× |
+ | 16 | ~10% | 3,071.5 | 715.8 | 0.3 | 1.4 | 4.29× |
+ | 16 | 100% | 2,667.7 | 715.1 | 0.4 | 1.4 | 3.73× |
+ | 256 | 0% | 2,823.3 | 1,270.9 | 0.4 | 0.8 | 2.22× |
+ | 256 | ~10% | 2,802.2 | 1,280.0 | 0.4 | 0.8 | 2.19× |
+ | 256 | 100% | 2,234.7 | 1,464.3 | 0.4 | 0.7 | 1.53× |
+ | 1,024 | 0% | 1,847.0 | 3,239.8 | 0.5 | 0.3 | 0.57× |
+ | 1,024 | ~10% | 1,808.6 | 3,287.6 | 0.6 | 0.3 | 0.55× |
+ | 1,024 | 100% | 1,195.0 | 4,004.1 | 0.8 | 0.2 | 0.30× |
+
+ These results led to the 256-byte default limit for replacement-column routing. Runs are independent single benchmark invocations, so treat small differences as directional rather than statistically conclusive.
## right
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index ddf3df2892..1bf7231a76 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -605,7 +605,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab
| `regexp_replace` | ✅ | Hybrid | |
| `regexp_substr` | ✅ | — | Rewrites to `nullif(regexp_extract(...), '')`; `regexp_extract` routes through the JVM codegen dispatcher by default |
| `repeat` | ✅ | Native | |
-| `replace` | ✅ | Hybrid | |
+| `replace` | ✅ | Hybrid | Native by default for a direct string column with a non-empty, valid UTF-8 search literal and either a replacement literal (both at most 64 KiB) or a replacement column (search at most 256 bytes and at most 8 MiB when broadcast per batch), under UTF8_BINARY collation; other cases use the JVM codegen dispatcher by default |
| `right` | ✅ | Native | |
| `rpad` | ✅ | Hybrid | String inputs use the native kernel with a column string and literal padding; literal strings and column padding use codegen dispatch. Binary inputs use codegen dispatch. |
| `rtrim` | ✅ | Hybrid | Non-UTF8_BINARY collated input with a trim string routes through the JVM codegen dispatcher; other input runs natively |
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 26ebad05a9..3ae5273584 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -19,7 +19,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, StringInstr, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTran [...]
+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, StringInstr, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringS [...]
import org.apache.spark.sql.types.{ArrayType, BinaryType, DataTypes, IntegerType, LongType, StringType}
import org.apache.spark.unsafe.types.UTF8String
@@ -186,13 +186,58 @@ object CometInitCap extends CometScalarFunction[InitCap]("initcap") with NativeO
object CometStringReplace
extends CometScalarFunction[StringReplace]("replace")
- with NativeOptInAvailable {
+ with NativeOptInAvailable
+ with CometTypeShim {
+
+ private val MaxNativeLiteralBytes = 64 * 1024
+ // With a replacement column, DataFusion broadcasts the scalar search into an Arrow array.
+ // Bound both the literal and its total expansion at the configured batch size.
+ private val MaxBroadcastSearchBytes = 256
+ private val MaxBroadcastSearchBytesPerBatch = 8 * 1024 * 1024
+
+ override def hasConditionalNativeDefault: Boolean = true
+
+ override def getCompatibleNotes(): Seq[String] = Seq(
+ "Native by default for a direct string column with a non-empty, valid UTF-8 search " +
+ "literal and either a replacement literal (both at most 64 KiB) or a replacement " +
+ "column (search at most 256 bytes and at most 8 MiB when broadcast per batch), under " +
+ "UTF8_BINARY collation. Other cases use " +
+ "the JVM codegen dispatcher by default.")
+
+ // DataFusion 55.1.0 keeps two non-null literal arguments as scalars when the source is an
+ // array. With a replacement column it broadcasts the search literal, hence the smaller bound.
+ // Direct column operands avoid Spark NULL short-circuit differences in nested expressions.
+ private def nativeSafeSubset(expr: StringReplace): Boolean = expr.children match {
+ case Seq(source, Literal(search: UTF8String, _), Literal(replacement: UTF8String, _)) =>
+ (source.isInstanceOf[Attribute] || source.isInstanceOf[BoundReference]) &&
+ search.numBytes() > 0 && search.numBytes() <= MaxNativeLiteralBytes &&
+ replacement.numBytes() <= MaxNativeLiteralBytes &&
+ isValidUtf8(search) && isValidUtf8(replacement) &&
+ !hasNonDefaultStringCollation(expr.dataType) &&
+ !expr.children.exists(child => hasNonDefaultStringCollation(child.dataType))
+ case Seq(source, Literal(search: UTF8String, _), replacement) =>
+ (source.isInstanceOf[Attribute] || source.isInstanceOf[BoundReference]) &&
+ (replacement.isInstanceOf[Attribute] || replacement.isInstanceOf[BoundReference]) &&
+ search.numBytes() > 0 && search.numBytes() <= MaxBroadcastSearchBytes &&
+ search.numBytes().toLong * CometConf.COMET_BATCH_SIZE.get() <=
+ MaxBroadcastSearchBytesPerBatch &&
+ isValidUtf8(search) &&
+ !hasNonDefaultStringCollation(expr.dataType) &&
+ !expr.children.exists(child => hasNonDefaultStringCollation(child.dataType))
+ case _ => false
+ }
+
+ private def useNative(expr: StringReplace): Boolean =
+ CometConf.isExprAllowIncompat(getExprConfigName(expr)) || nativeSafeSubset(expr)
override def getIncompatibleReasons(): Seq[String] =
- Seq("Produces different results from Spark when the search string is empty")
+ Seq(
+ "Unrestricted native arguments can change malformed UTF-8 results or Spark's NULL " +
+ "short-circuit behavior, ignore non-default collations, or overflow Arrow offsets when " +
+ "scalars are broadcast")
override def getSupportLevel(expr: StringReplace): SupportLevel =
- if (!CometConf.isExprAllowIncompat(getExprConfigName(expr))) {
+ if (!useNative(expr)) {
Compatible(nativeOptIn =
Some(NativeOptIn(CometConf.getExprAllowIncompatConfigKey(getExprConfigName(expr)))))
} else {
@@ -203,10 +248,7 @@ object CometStringReplace
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 (useNative(expr)) {
super.convert(expr, inputs, binding)
} else {
// Run Spark's own generated code inside the Comet pipeline so the result matches Spark
diff --git a/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_disabled.sql b/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_disabled.sql
index 96357a6b96..61c2a8b918 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_disabled.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_disabled.sql
@@ -29,8 +29,13 @@ INSERT INTO routing_string_opt_in VALUES ('hello world', '{"a":1}'), ('', '{}'),
query expect_fallback(initcap: spark.comet.exec.scalaUDF.codegen.enabled=false)
SELECT initcap(s) FROM routing_string_opt_in
-query expect_fallback(replace: spark.comet.exec.scalaUDF.codegen.enabled=false)
+-- A direct column with bounded, non-empty literals is native even without opt-in.
+query expect_native(replace)
SELECT replace(s, 'l', 'L') FROM routing_string_opt_in
+-- Empty search remains outside the native default subset and needs the dispatcher.
+query expect_fallback(replace: spark.comet.exec.scalaUDF.codegen.enabled=false)
+SELECT replace(s, '', 'L') FROM routing_string_opt_in
+
query expect_fallback(get_json_object: spark.comet.exec.scalaUDF.codegen.enabled=false)
SELECT get_json_object(j, '$.a') FROM routing_string_opt_in
diff --git a/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_enabled.sql b/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_enabled.sql
index ca2f091c05..b851d7991c 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_enabled.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/routing_string_opt_in_enabled.sql
@@ -29,8 +29,13 @@ INSERT INTO routing_string_opt_in VALUES ('hello world', '{"a":1}'), ('', '{}'),
query expect_dispatch(initcap)
SELECT initcap(s) FROM routing_string_opt_in
-query expect_dispatch(replace)
+-- A direct column with bounded, non-empty literals is native even without opt-in.
+query expect_native(replace)
SELECT replace(s, 'l', 'L') FROM routing_string_opt_in
+-- Empty search remains outside the native default subset and needs the dispatcher.
+query expect_dispatch(replace)
+SELECT replace(s, '', 'L') FROM routing_string_opt_in
+
query expect_dispatch(get_json_object)
SELECT get_json_object(j, '$.a') FROM routing_string_opt_in
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 db4800ca57..80b3a89ec9 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
@@ -24,10 +24,7 @@ INSERT INTO test_str_replace VALUES ('hello world', 'world', 'there'), ('aaa', '
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.
+-- Empty literal search: CometStringReplace retains the codegen dispatcher path.
-- https://github.com/apache/datafusion-comet/issues/4497
query
SELECT replace('hello', '', 'x')
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
index 1c1cbf164f..6bdad65ebf 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -855,6 +855,150 @@ class CometCodegenSuite
}
}
+ test("replace with safe literal arguments runs natively by default") {
+ 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,
+ CometConf.getExprAllowIncompatConfigKey("StringReplace") -> "false") {
+ withTable("t") {
+ sql("CREATE TABLE t (s STRING) USING parquet")
+ sql("INSERT INTO t VALUES ('aaaa'), ('你好你好'), (''), (NULL)")
+
+ Seq(
+ "SELECT replace(s, 'aa', 'x') FROM t",
+ "SELECT replace(s, '你好', '世界') FROM t",
+ "SELECT replace(s, 'aa', 'aaa') FROM t",
+ "SELECT replace(s, 'missing', '') FROM t").foreach { query =>
+ val df = sql(query)
+ assertReplaceDispatch(df, expectDispatcher = false, query)
+ val plan = df.queryExecution.executedPlan
+ val info = new ExtendedExplainInfo()
+ assert(
+ info.getNativeExpressions(plan).contains("replace"),
+ s"expected native replace for $query:\n${info.generateExtendedInfo(plan)}")
+ }
+ }
+ }
+ }
+
+ test("replace literal arguments route at the 64 KiB boundary") {
+ 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,
+ CometConf.getExprAllowIncompatConfigKey("StringReplace") -> "false") {
+ withTable("t") {
+ sql("CREATE TABLE t (s STRING) USING parquet")
+ sql("INSERT INTO t VALUES ('a'), ('other'), (NULL), (repeat('b', 65536))")
+
+ Seq(
+ ("SELECT replace(s, repeat('b', 65536), 'x') FROM t", false),
+ ("SELECT replace(s, repeat('b', 65537), 'x') FROM t", true),
+ ("SELECT replace(s, 'a', repeat('x', 65536)) FROM t", false),
+ ("SELECT replace(s, 'a', repeat('x', 65537)) FROM t", true)).foreach {
+ case (query, expectDispatcher) =>
+ val df = sql(query)
+ assertReplaceDispatch(df, expectDispatcher, query)
+ if (!expectDispatcher) {
+ val plan = df.queryExecution.executedPlan
+ val info = new ExtendedExplainInfo()
+ assert(
+ info.getNativeExpressions(plan).contains("replace"),
+ s"expected native replace for $query:\n${info.generateExtendedInfo(plan)}")
+ }
+ }
+ }
+ }
+ }
+
+ test("replace with a literal search and replacement column runs natively by default") {
+ 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,
+ CometConf.getExprAllowIncompatConfigKey("StringReplace") -> "false") {
+ withTable("t") {
+ sql("CREATE TABLE t (s STRING, r STRING, search STRING) USING parquet")
+ sql("""INSERT INTO t VALUES
+ |('aaaa', 'x', 'aa'), ('你好你好', '世界', '你好'),
+ |('nomatch', 'z', 'm'), ('aaaa', '', ''),
+ |('', 'x', ''), (NULL, 'q', 'a'), ('aaa', NULL, 'a')""".stripMargin)
+
+ Seq(
+ "SELECT replace(s, 'aa', r) FROM t",
+ "SELECT replace(s, '你好', r) FROM t",
+ "SELECT replace(s, repeat('a', 256), r) FROM t").foreach { query =>
+ val df = sql(query)
+ assertReplaceDispatch(df, expectDispatcher = false, query)
+ val plan = df.queryExecution.executedPlan
+ val info = new ExtendedExplainInfo()
+ assert(
+ info.getNativeExpressions(plan).contains("replace"),
+ s"expected native replace for $query:\n${info.generateExtendedInfo(plan)}")
+ }
+
+ Seq(
+ "SELECT replace(s, '', r) FROM t",
+ "SELECT replace(s, search, r) FROM t",
+ "SELECT replace(s, repeat('a', 257), r) FROM t",
+ "SELECT replace(s, repeat('a', 1024), r) FROM t").foreach { query =>
+ assertReplaceDispatch(sql(query), expectDispatcher = true, query)
+ }
+
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "65536") {
+ val query = "SELECT replace(s, repeat('a', 256), r) FROM t"
+ assertReplaceDispatch(sql(query), expectDispatcher = true, query)
+ }
+
+ if (isSpark41Plus) {
+ val query = "SELECT replace(s COLLATE UTF8_LCASE, 'aa', r) FROM t"
+ assertReplaceDispatch(sql(query), expectDispatcher = true, query)
+ }
+ }
+ }
+ }
+
+ test("replace broadcast limit routes actual large batches at 8 MiB") {
+ 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,
+ CometConf.getExprAllowIncompatConfigKey("StringReplace") -> "false") {
+ withTempPath { dir =>
+ spark
+ .range(0, 32769, 1, numPartitions = 1)
+ .selectExpr("repeat('a', 256) AS s", "CASE WHEN id % 2 = 0 THEN 'x' ELSE 'yz' END AS r")
+ .write
+ .parquet(dir.getCanonicalPath)
+ withTable("t") {
+ sql(s"CREATE TABLE t USING parquet LOCATION '${dir.getCanonicalPath}'")
+ val query = "SELECT replace(s, repeat('a', 256), r) FROM t"
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "32768") {
+ val df = sql(query)
+ assertReplaceDispatch(df, expectDispatcher = false, query)
+ val plan = df.queryExecution.executedPlan
+ val info = new ExtendedExplainInfo()
+ assert(
+ info.getNativeExpressions(plan).contains("replace"),
+ s"expected native replace for $query:\n${info.generateExtendedInfo(plan)}")
+ }
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "32769") {
+ assertReplaceDispatch(sql(query), expectDispatcher = true, query)
+ }
+ }
+ }
+ }
+ }
+
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.
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
index aede98a129..e28afe3374 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
@@ -41,7 +41,11 @@ case class StringExprConfig(
* {{{
* SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometStringExpressionBenchmark
* }}}
- * Results will be written to "spark/benchmarks/CometStringExpressionBenchmark-**results.txt".
+ * Results will be written to "spark/benchmarks/CometStringExpressionBenchmark-**results.txt". The
+ * replace_match cases use literal arguments, while replace_column_match cases read the
+ * replacement from a column. Each set uses the same input width with 0%, about 10%, and 100%
+ * match density. Compare their Comet results on a dispatcher-only revision and a native-default
+ * revision under the same Spark profile, batch size, and hardware.
*/
object CometStringExpressionBenchmark extends CometBenchmarkBase {
@@ -72,6 +76,25 @@ object CometStringExpressionBenchmark extends CometBenchmarkBase {
"select regexp_replace(c1, '[0-9]', 'X') from parquetV1Table"),
StringExprConfig("repeat", "select repeat(c1, 3) from parquetV1Table"),
StringExprConfig("replace", "select replace(c1, '123', 'ab') from parquetV1Table"),
+ // Both literal arguments stay scalar in DataFusion 55.1.0's replace fast path.
+ StringExprConfig(
+ "replace_match_0pct",
+ "select replace(replace_no_match, '123', 'ab') from parquetV1Table"),
+ StringExprConfig(
+ "replace_match_10pct",
+ "select replace(replace_ten_percent, '123', 'ab') from parquetV1Table"),
+ StringExprConfig(
+ "replace_match_100pct",
+ "select replace(replace_all_match, '123', 'ab') from parquetV1Table"),
+ StringExprConfig(
+ "replace_column_match_0pct",
+ "select replace(replace_no_match, '123', replace_with) from parquetV1Table"),
+ StringExprConfig(
+ "replace_column_match_10pct",
+ "select replace(replace_ten_percent, '123', replace_with) from parquetV1Table"),
+ StringExprConfig(
+ "replace_column_match_100pct",
+ "select replace(replace_all_match, '123', replace_with) from parquetV1Table"),
StringExprConfig("reverse", "select reverse(c1) from parquetV1Table"),
StringExprConfig("rlike", "select c1 rlike '[0-9]+' from parquetV1Table"),
StringExprConfig("rpad", "select rpad(c1, 150, 'x') from parquetV1Table"),
@@ -85,7 +108,23 @@ object CometStringExpressionBenchmark extends CometBenchmarkBase {
StringExprConfig("substring", "select substring(c1, 1, 100) from parquetV1Table"),
StringExprConfig("translate", "select translate(c1, '123456', 'aBcDeF') from parquetV1Table"),
StringExprConfig("trim", "select trim(c1) from parquetV1Table"),
- StringExprConfig("upper", "select upper(c1) from parquetV1Table"))
+ StringExprConfig("upper", "select upper(c1) from parquetV1Table")) ++
+ Seq(16, 256, 1024).flatMap { searchBytes =>
+ val search = "'" + ("a" * searchBytes) + "'"
+ Seq(
+ StringExprConfig(
+ s"replace_column_search_${searchBytes}b_0pct",
+ s"select replace(replace_long_${searchBytes}_0pct, $search, replace_with) " +
+ "from parquetV1Table"),
+ StringExprConfig(
+ s"replace_column_search_${searchBytes}b_10pct",
+ s"select replace(replace_long_${searchBytes}_10pct, $search, replace_with) " +
+ "from parquetV1Table"),
+ StringExprConfig(
+ s"replace_column_search_${searchBytes}b_100pct",
+ s"select replace(replace_long_${searchBytes}_100pct, $search, replace_with) " +
+ "from parquetV1Table"))
+ }
// Collated cases are Spark 4.0+ only, since the COLLATE syntax does not parse on 3.4 and 3.5.
// CometLevenshtein reports collated input as Unsupported and CodegenDispatchFallback runs it
@@ -107,24 +146,60 @@ object CometStringExpressionBenchmark extends CometBenchmarkBase {
}
override def runCometBenchmark(mainArgs: Array[String]): Unit = {
- runBenchmarkWithTable("String expressions", 1024) { v =>
+ val replaceMatchOnly = mainArgs.contains("replace-match-only")
+ val replaceColumnOnly = mainArgs.contains("replace-column-only")
+ val replaceLongSearchOnly = mainArgs.contains("replace-column-long-only")
+ val rows =
+ if (replaceMatchOnly || replaceColumnOnly || replaceLongSearchOnly) 65536 else 1024
+ runBenchmarkWithTable("String expressions", rows) { v =>
withTempPath { dir =>
withTempTable("parquetV1Table") {
// c2 gives expressions that take a length or a count something to vary over per row.
// `pmod` keeps it non-negative, so `space(c2)` builds a string on every row rather
// than returning empty for the negative half of the input.
+ val longSearchInputs = Seq(16, 256, 1024).flatMap { searchBytes =>
+ val flankBytes = (2048 - searchBytes) / 2
+ val matchSource =
+ s"CONCAT(REPEAT('b', $flankBytes), REPEAT('a', $searchBytes), " +
+ s"REPEAT('b', $flankBytes))"
+ Seq(
+ s"REPEAT('b', 2048) AS replace_long_${searchBytes}_0pct",
+ s"CASE WHEN PMOD(value, 10) = 0 THEN $matchSource ELSE " +
+ s"REPEAT('b', 2048) END AS replace_long_${searchBytes}_10pct",
+ s"$matchSource AS replace_long_${searchBytes}_100pct")
+ }
prepareTable(
dir,
spark.sql(
"SELECT REPEAT(CAST(value AS STRING), 10) AS c1," +
- s" CAST(PMOD(value, 200) AS INT) AS c2 FROM $tbl"))
+ " CAST(PMOD(value, 200) AS INT) AS c2," +
+ " REPEAT('abc---def', 8) AS replace_no_match," +
+ " CASE WHEN PMOD(value, 10) = 0 THEN REPEAT('abc123def', 8)" +
+ " ELSE REPEAT('abc---def', 8) END AS replace_ten_percent," +
+ " REPEAT('abc123def', 8) AS replace_all_match," +
+ " CASE WHEN PMOD(value, 20) = 0 THEN CAST(NULL AS STRING)" +
+ " WHEN PMOD(value, 3) = 0 THEN ''" +
+ " WHEN PMOD(value, 3) = 1 THEN 'ab' ELSE 'XYZ' END AS replace_with" +
+ ", " + longSearchInputs.mkString(", ") +
+ s" FROM $tbl"))
val extraConfigs = Map(
CometConf.getExprAllowIncompatConfigKey("Upper") -> "true",
CometConf.getExprAllowIncompatConfigKey("Lower") -> "true",
CometConf.getExprAllowIncompatConfigKey("InitCap") -> "true")
- (stringExpressions ++ collatedStringExpressions).foreach { config =>
+ val cases = if (replaceColumnOnly) {
+ stringExpressions.filter(_.name.startsWith("replace_column_match_"))
+ } else if (replaceLongSearchOnly) {
+ stringExpressions.filter(_.name.startsWith("replace_column_search_"))
+ } else if (replaceMatchOnly) {
+ stringExpressions.filter(config =>
+ config.name.startsWith("replace_match_") ||
+ config.name.startsWith("replace_column_match_"))
+ } else {
+ stringExpressions ++ collatedStringExpressions
+ }
+ cases.foreach { config =>
val allConfigs = extraConfigs ++ config.extraCometConfigs
runBenchmark(config.name) {
runExpressionBenchmark(config.name, v.toLong, config.query, allConfigs)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org