You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/09/08 21:11:58 UTC
(datafusion-comet) branch main updated: feat: enable codegen dispatch for `lpad` and `rpad` (#5764)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 49a4b6c526 feat: enable codegen dispatch for `lpad` and `rpad` (#5764)
49a4b6c526 is described below
commit 49a4b6c5268b8d4dc8306ea8a2302f4d0c95ffbb
Author: Daipayan Mukherjee <da...@gmail.com>
AuthorDate: Tue Sep 8 22:11:53 2026 +0100
feat: enable codegen dispatch for `lpad` and `rpad` (#5764)
* feat: route unsupported string padding through codegen dispatch
* style: format the padding benchmark
---
docs/source/user-guide/latest/expressions.md | 4 +-
.../main/scala/org/apache/comet/GenerateDocs.scala | 11 +++-
.../scala/org/apache/comet/serde/strings.scala | 4 +-
.../sql-tests/expressions/string/string_lpad.sql | 19 +++++--
.../{string_rpad.sql => string_lpad_fallback.sql} | 28 +++++-----
.../sql-tests/expressions/string/string_rpad.sql | 19 +++++--
.../{string_rpad.sql => string_rpad_fallback.sql} | 28 +++++-----
.../apache/comet/CometStringExpressionSuite.scala | 59 +++++++++++++++++++---
.../CometPaddingExpressionBenchmark.scala | 59 ++++++++++++++++++++++
9 files changed, 183 insertions(+), 48 deletions(-)
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 1fb979fd3c..01d898cd46 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -579,7 +579,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
| `levenshtein` | ✅ | Hybrid | Non-UTF8_BINARY collated input routes through the JVM codegen dispatcher; other input runs natively |
| `locate` | ✅ | Codegen dispatch | |
| `lower` | ✅ | Hybrid | |
-| `lpad` | ✅ | — | |
+| `lpad` | ✅ | 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. |
| `ltrim` | ✅ | Native | |
| `luhn_check` | ✅ | — | Native via `StaticInvoke` (tests: luhn_check.sql) |
| `mask` | ✅ | — | Routed through the JVM codegen dispatcher |
@@ -596,7 +596,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
| `repeat` | ✅ | Native | |
| `replace` | ✅ | Hybrid | |
| `right` | ✅ | Native | |
-| `rpad` | ✅ | — | |
+| `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` | ✅ | Native | |
| `soundex` | ✅ | Native | |
| `space` | ✅ | Native | |
diff --git a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala
index e5d0c94215..b4bc6c3b43 100644
--- a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala
+++ b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala
@@ -25,7 +25,7 @@ import scala.collection.mutable
import scala.collection.mutable.ListBuffer
import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
-import org.apache.spark.sql.catalyst.expressions.{Cast, Expression}
+import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, StringLPad, StringRPad}
import org.apache.comet.CometConf.COMET_ONHEAP_MEMORY_OVERHEAD
import org.apache.comet.expressions.{CometCast, CometEvalMode}
@@ -198,7 +198,14 @@ object GenerateDocs {
(exprKinds ++ aggKinds).toMap
}
FunctionRegistry.expressions.iterator.map { case (name, (info, _)) =>
- name -> classNameToKind.getOrElse(info.getClassName, ImplUnknown)
+ // Spark registers builders for padding because the expression depends on the input type.
+ // Resolve their string expressions, which expose both the native and dispatcher paths.
+ val className = name match {
+ case "lpad" => classOf[StringLPad].getName
+ case "rpad" => classOf[StringRPad].getName
+ case _ => info.getClassName
+ }
+ name -> classNameToKind.getOrElse(className, ImplUnknown)
}.toMap
}
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 4665c95461..a5707c151f 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -421,7 +421,7 @@ private object PadReasons {
val nonLiteralPadReason = "Only scalar values are supported for the `pad` argument."
}
-object CometStringRPad extends CometExpressionSerde[StringRPad] {
+object CometStringRPad extends CometExpressionSerde[StringRPad] with CodegenDispatchFallback {
override def getUnsupportedReasons(): Seq[String] =
Seq(PadReasons.literalStrReason, PadReasons.nonLiteralPadReason)
@@ -449,7 +449,7 @@ object CometStringRPad extends CometExpressionSerde[StringRPad] {
}
}
-object CometStringLPad extends CometExpressionSerde[StringLPad] {
+object CometStringLPad extends CometExpressionSerde[StringLPad] with CodegenDispatchFallback {
override def getUnsupportedReasons(): Seq[String] =
Seq(PadReasons.literalStrReason, PadReasons.nonLiteralPadReason)
diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql
index 93867576aa..27b38bbda5 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql
@@ -19,11 +19,24 @@ statement
CREATE TABLE test_lpad(s string, len int, pad string) USING parquet
statement
-INSERT INTO test_lpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x')
+INSERT INTO test_lpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x'), ('hi', 5, NULL), ('hi', 5, ''), ('', 3, ''), ('hi', -100, 'x'), (NULL, NULL, NULL), ('"hi', 7, '"x'), ('café', 7, '中文'), ('é', 5, '🙂')
-query expect_fallback(Only scalar values are supported for the `pad` argument)
+-- Column padding runs through the codegen dispatcher (issue #5579).
+query
SELECT lpad(s, len, pad) FROM test_lpad
+query
+SELECT lpad(s, 5, pad) FROM test_lpad
+
+query
+SELECT lpad('hi', len, pad) FROM test_lpad
+
+query
+SELECT lpad('hi', len, 'xy') FROM test_lpad
+
+query
+SELECT lpad('hi', len) FROM test_lpad
+
query
SELECT lpad(s, len) FROM test_lpad
@@ -36,5 +49,5 @@ query
SELECT lpad(s, 5, 'x') FROM test_lpad
-- literal + literal + literal
-query expect_fallback(Scalar values are not supported for the `str` argument)
+query
SELECT lpad('hi', 5, 'x'), lpad('hello', 3, 'x'), lpad('', 3, 'a'), lpad(NULL, 5, 'x')
diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_lpad_fallback.sql
similarity index 53%
copy from spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
copy to spark/src/test/resources/sql-tests/expressions/string/string_lpad_fallback.sql
index 5832f48785..32c476379d 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/string_lpad_fallback.sql
@@ -15,26 +15,26 @@
-- specific language governing permissions and limitations
-- under the License.
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=false
+
statement
-CREATE TABLE test_rpad(s string, len int, pad string) USING parquet
+CREATE TABLE test_lpad_fallback(s string, len int, pad string) USING parquet
statement
-INSERT INTO test_rpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x')
+INSERT INTO test_lpad_fallback VALUES ('hi', 5, 'xy'), ('hello', 3, 'x'), (NULL, NULL, NULL)
-query expect_fallback(Only scalar values are supported for the `pad` argument)
-SELECT rpad(s, len, pad) FROM test_rpad
+query expect_fallback(spark.comet.exec.scalaUDF.codegen.enabled)
+SELECT lpad(s, len, pad) FROM test_lpad_fallback
-query
-SELECT rpad(s, len) FROM test_rpad
+query expect_fallback(spark.comet.exec.scalaUDF.codegen.enabled)
+SELECT lpad('hi', len, 'xy') FROM test_lpad_fallback
--- column + column + literal
-query
-SELECT rpad(s, len, 'x') FROM test_rpad
+query expect_fallback(spark.comet.exec.scalaUDF.codegen.enabled)
+SELECT lpad('hi', 5, 'xy')
--- column + literal + literal
+-- The native argument shapes do not require the dispatcher.
query
-SELECT rpad(s, 5, 'x') FROM test_rpad
+SELECT lpad(s, len, 'xy') FROM test_lpad_fallback
--- literal + literal + literal
-query expect_fallback(Scalar values are not supported for the `str` argument)
-SELECT rpad('hi', 5, 'x'), rpad('hello', 3, 'x'), rpad('', 3, 'a'), rpad(NULL, 5, 'x')
+query
+SELECT lpad(s, len) FROM test_lpad_fallback
diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
index 5832f48785..a5f9951273 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
@@ -19,11 +19,24 @@ statement
CREATE TABLE test_rpad(s string, len int, pad string) USING parquet
statement
-INSERT INTO test_rpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x')
+INSERT INTO test_rpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x'), ('hi', 5, NULL), ('hi', 5, ''), ('', 3, ''), ('hi', -100, 'x'), (NULL, NULL, NULL), ('"hi', 7, '"x'), ('café', 7, '中文'), ('é', 5, '🙂')
-query expect_fallback(Only scalar values are supported for the `pad` argument)
+-- Column padding runs through the codegen dispatcher (issue #5579).
+query
SELECT rpad(s, len, pad) FROM test_rpad
+query
+SELECT rpad(s, 5, pad) FROM test_rpad
+
+query
+SELECT rpad('hi', len, pad) FROM test_rpad
+
+query
+SELECT rpad('hi', len, 'xy') FROM test_rpad
+
+query
+SELECT rpad('hi', len) FROM test_rpad
+
query
SELECT rpad(s, len) FROM test_rpad
@@ -36,5 +49,5 @@ query
SELECT rpad(s, 5, 'x') FROM test_rpad
-- literal + literal + literal
-query expect_fallback(Scalar values are not supported for the `str` argument)
+query
SELECT rpad('hi', 5, 'x'), rpad('hello', 3, 'x'), rpad('', 3, 'a'), rpad(NULL, 5, 'x')
diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_rpad_fallback.sql
similarity index 53%
copy from spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
copy to spark/src/test/resources/sql-tests/expressions/string/string_rpad_fallback.sql
index 5832f48785..8ac5bb36ba 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/string_rpad_fallback.sql
@@ -15,26 +15,26 @@
-- specific language governing permissions and limitations
-- under the License.
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=false
+
statement
-CREATE TABLE test_rpad(s string, len int, pad string) USING parquet
+CREATE TABLE test_rpad_fallback(s string, len int, pad string) USING parquet
statement
-INSERT INTO test_rpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x')
+INSERT INTO test_rpad_fallback VALUES ('hi', 5, 'xy'), ('hello', 3, 'x'), (NULL, NULL, NULL)
-query expect_fallback(Only scalar values are supported for the `pad` argument)
-SELECT rpad(s, len, pad) FROM test_rpad
+query expect_fallback(spark.comet.exec.scalaUDF.codegen.enabled)
+SELECT rpad(s, len, pad) FROM test_rpad_fallback
-query
-SELECT rpad(s, len) FROM test_rpad
+query expect_fallback(spark.comet.exec.scalaUDF.codegen.enabled)
+SELECT rpad('hi', len, 'xy') FROM test_rpad_fallback
--- column + column + literal
-query
-SELECT rpad(s, len, 'x') FROM test_rpad
+query expect_fallback(spark.comet.exec.scalaUDF.codegen.enabled)
+SELECT rpad('hi', 5, 'xy')
--- column + literal + literal
+-- The native argument shapes do not require the dispatcher.
query
-SELECT rpad(s, 5, 'x') FROM test_rpad
+SELECT rpad(s, len, 'xy') FROM test_rpad_fallback
--- literal + literal + literal
-query expect_fallback(Scalar values are not supported for the `str` argument)
-SELECT rpad('hi', 5, 'x'), rpad('hello', 3, 'x'), rpad('', 3, 'a'), rpad(NULL, 5, 'x')
+query
+SELECT rpad(s, len) FROM test_rpad_fallback
diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
index e67125e389..c58403487a 100644
--- a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
@@ -46,6 +46,53 @@ class CometStringExpressionSuite extends CometTestBase with CometCodegenAssertio
testStringPadding("rpad")
}
+ for ((function, expressionName) <- Seq("lpad" -> "StringLPad", "rpad" -> "StringRPad")) {
+ test(s"$function dispatches unsupported argument shapes (issue #5579)") {
+ val data: Seq[(String, Option[Int], String)] = Seq(
+ ("hi", Some(5), "xy"),
+ ("hello", Some(3), "x"),
+ ("", Some(3), "a"),
+ ("hi", Some(5), ""),
+ (null, Some(5), "x"),
+ ("hi", None, "x"),
+ ("hi", Some(5), null),
+ (null, None, null))
+ withParquetTable(data, "tbl") {
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") {
+ for (allowIncompatible <- Seq("false", "true")) {
+ withSQLConf(
+ CometConf.getExprAllowIncompatConfigKey(expressionName) -> allowIncompatible) {
+ for (query <- Seq(
+ s"SELECT $function(_1, _2, _3) FROM tbl",
+ s"SELECT $function('hi', _2, 'xy') FROM tbl",
+ s"SELECT $function('hi', 5, 'xy') FROM tbl")) {
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(query)
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test(s"$function keeps supported argument shapes native") {
+ withParquetTable(Seq(("hi", 5), ("hello", 3), ("", 0)), "tbl") {
+ for (query <- Seq(
+ s"SELECT $function(_1, _2) FROM tbl",
+ s"SELECT $function(_1, _2, 'xy') FROM tbl")) {
+ CometScalaUDFCodegen.resetStats()
+ checkSparkAnswerAndOperator(query)
+ assert(
+ CometScalaUDFCodegen.stats().totalLookups == 0,
+ s"expected native execution for $query")
+ }
+ }
+ }
+ }
+
test("lpad/rpad with NULL length") {
// FuzzDataGenerator never generates NULL integers (#5389), so build the rows explicitly.
// Spark's StringLPad/StringRPad are null-intolerant: a NULL length yields a NULL row.
@@ -109,14 +156,10 @@ class CometStringExpressionSuite extends CometTestBase with CometCodegenAssertio
// all arguments are literal, so Spark constant folding will kick in
// and pad function will not be evaluated by Comet
checkSparkAnswerAndOperator(sql)
- } else if (isLiteralStr) {
- checkSparkAnswerAndFallbackReason(
- sql,
- "Scalar values are not supported for the `str` argument")
- } else if (!isLiteralPad) {
- checkSparkAnswerAndFallbackReason(
- sql,
- "Only scalar values are supported for the `pad` argument")
+ } else if (isLiteralStr || !isLiteralPad) {
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql)
+ }
} else {
checkSparkAnswerAndOperator(sql)
}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPaddingExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPaddingExpressionBenchmark.scala
new file mode 100644
index 0000000000..6fab1ea7a7
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPaddingExpressionBenchmark.scala
@@ -0,0 +1,59 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.benchmark
+
+/**
+ * Compares padding through codegen dispatch with Spark and the existing native argument shape.
+ * Run with:
+ * {{{
+ * make benchmark-org.apache.spark.sql.benchmark.CometPaddingExpressionBenchmark
+ * }}}
+ */
+object CometPaddingExpressionBenchmark extends CometBenchmarkBase {
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ val rows = 1024 * 1024
+ withTempPath { dir =>
+ withTempTable("parquetV1Table") {
+ prepareTable(
+ dir,
+ spark
+ .range(rows)
+ .selectExpr(
+ "CAST(id AS STRING) AS s",
+ "CAST(id % 32 + 8 AS INT) AS len",
+ "CASE WHEN id % 2 = 0 THEN 'xy' ELSE 'z' END AS pad"))
+ for (function <- Seq("lpad", "rpad")) {
+ for ((shape, arguments) <- Seq(
+ "column padding" -> "s, len, pad",
+ "literal string" -> "'hi', len, 'xy'",
+ "native" -> "s, len, 'xy'")) {
+ val name = s"$function ($shape)"
+ runBenchmark(name) {
+ runExpressionBenchmark(
+ name,
+ rows,
+ s"SELECT $function($arguments) FROM parquetV1Table")
+ }
+ }
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org