This is an automated email from the ASF dual-hosted git repository.
mbutrovich 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 83339d8891 feat: hint at native pyarrow UDF path when the feature is
disabled (#4892)
83339d8891 is described below
commit 83339d88914ed1e588f541882f3b6ddc27819325
Author: Andy Grove <[email protected]>
AuthorDate: Fri Jul 17 09:18:11 2026 -0600
feat: hint at native pyarrow UDF path when the feature is disabled (#4892)
* feat: hint at native pyarrow UDF path when the feature is disabled
When spark.comet.exec.pyarrowUdf.enabled is off (the default), mapInArrow
and mapInPandas operators silently fall back to Spark. Annotate the operator
with a non-fallback [COMET-INFO] hint naming the config to opt into the
native CometMapInBatchExec path, reusing the existing NativeOptIn wording.
The hint fires only on operators that would actually run natively if the
flag were flipped: the eligibility check is unchanged, the conf gate simply
moves from the extractor into the rewrite arm.
* docs: remove emphasis from EliminateRedundantTransitions comment
---
.../rules/EliminateRedundantTransitions.scala | 49 ++++++++++++++--------
.../spark/sql/comet/CometMapInBatchSuite.scala | 23 +++++++++-
2 files changed, 54 insertions(+), 18 deletions(-)
diff --git
a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
index 477c2afd99..4b4b8c13bd 100644
---
a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
+++
b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
@@ -30,6 +30,8 @@ import org.apache.spark.sql.execution.adaptive.QueryStageExec
import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.withInfo
+import org.apache.comet.serde.NativeOptIn
import org.apache.comet.shims.ShimSQLConf
// This rule is responsible for eliminating redundant transitions between
row-based and
@@ -114,13 +116,26 @@ case class EliminateRedundantTransitions(session:
SparkSession)
// source string/binary vectors always use 4-byte offsets while the
destination root is
// allocated with 8-byte offsets when this conf is on. The buffer counts
match but the
// offset width does not, so a direct memcpy would corrupt the offsets.
- case EligibleMapInBatch(info, columnarChild) =>
- CometMapInBatchExec(
- info.func,
- info.output,
- columnarChild,
- info.isBarrier,
- info.pythonEvalType)
+ //
+ // `EligibleMapInBatch` matches whenever the operator would run natively
if the feature were
+ // enabled. When it is disabled (the default) we leave the vanilla Spark
operator in place
+ // but annotate it with a non-fallback `[COMET-INFO]` hint so the user
knows the native path
+ // exists behind a config flag.
+ case p @ EligibleMapInBatch(info, columnarChild) =>
+ if (CometConf.COMET_PYARROW_UDF_ENABLED.get()) {
+ CometMapInBatchExec(
+ info.func,
+ info.output,
+ columnarChild,
+ info.isBarrier,
+ info.pythonEvalType)
+ } else {
+ withInfo(
+ p,
+ NativeOptIn.message(
+ "PyArrow UDFs (mapInArrow/mapInPandas)",
+ CometConf.COMET_PYARROW_UDF_ENABLED.key))
+ }
// Spark adds `RowToColumnar` under Comet columnar shuffle. But it's
redundant as the
// shuffle takes row-based input.
@@ -174,19 +189,19 @@ case class EliminateRedundantTransitions(session:
SparkSession)
}
/**
- * Matches the plans this rule should rewrite to `CometMapInBatchExec`.
Single extractor used in
- * the `transformUp` arm above so the matchers and conf reads run once per
visited plan. Returns
- * `(info, columnarChild)` where `columnarChild` is the Comet columnar
producer that
- * `CometMapInBatchExec` will consume directly. Returns `None` (and the arm
misses) when the
- * conf is off, when `useLargeVarTypes` forces the fallback, when the plan
is not one of the
- * version-shimmed MapInArrow / MapInPandas operators, or when the child is
not a Comet
- * columnar-to-row transition we can strip.
+ * Matches the plans that could run natively as `CometMapInBatchExec`,
independent of whether
+ * the `spark.comet.exec.pyarrowUdf.enabled` feature flag is set. The
`transformUp` arm reads
+ * that flag to decide between rewriting the operator and merely annotating
it with an opt-in
+ * hint. Single extractor so the matchers run once per visited plan. Returns
`(info,
+ * columnarChild)` where `columnarChild` is the Comet columnar producer that
+ * `CometMapInBatchExec` will consume directly. Returns `None` (and the arm
misses) when
+ * `useLargeVarTypes` forces the fallback, when the plan is not one of the
version-shimmed
+ * MapInArrow / MapInPandas operators, or when the child is not a Comet
columnar-to-row
+ * transition we can strip.
*/
private object EligibleMapInBatch {
def unapply(plan: SparkPlan): Option[(MapInBatchInfo, SparkPlan)] = {
- if (!CometConf.COMET_PYARROW_UDF_ENABLED.get()) {
- None
- } else if (arrowUseLargeVarTypes(plan.conf)) {
+ if (arrowUseLargeVarTypes(plan.conf)) {
None
} else {
matchMapInArrow(plan)
diff --git
a/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala
b/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala
index c9acd2cc3f..09802dddaa 100644
---
a/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala
+++
b/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala
@@ -30,7 +30,7 @@ import org.apache.spark.sql.execution.python.MapInArrowExec
import org.apache.spark.sql.types.{LongType, StructField, StructType}
import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.comet.CometConf
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
import org.apache.comet.rules.EliminateRedundantTransitions
/** Minimal CometPlan leaf used to anchor the rule's transform without
triggering execution. */
@@ -105,6 +105,27 @@ class CometMapInBatchSuite extends CometTestBase {
}
}
+ test("rule annotates operator with opt-in hint when feature is disabled") {
+ withSQLConf(CometConf.COMET_PYARROW_UDF_ENABLED.key -> "false") {
+ val rewritten = EliminateRedundantTransitions(spark).apply(buildPlan())
+ val info = new ExtendedExplainInfo().generateExtendedInfo(rewritten)
+ assert(info.contains("[COMET-INFO:"), s"expected a [COMET-INFO: hint
in:\n$info")
+ assert(
+ info.contains(CometConf.COMET_PYARROW_UDF_ENABLED.key),
+ s"expected the opt-in config key in the hint:\n$info")
+ }
+ }
+
+ test("rule emits no opt-in hint when feature is enabled") {
+ withSQLConf(CometConf.COMET_PYARROW_UDF_ENABLED.key -> "true") {
+ val rewritten = EliminateRedundantTransitions(spark).apply(buildPlan())
+ val info = new ExtendedExplainInfo().generateExtendedInfo(rewritten)
+ assert(
+ !info.contains(CometConf.COMET_PYARROW_UDF_ENABLED.key),
+ s"unexpected opt-in hint when feature is enabled:\n$info")
+ }
+ }
+
test("rule rewrites chained MapInArrowExec into stacked
CometMapInBatchExec") {
// df.mapInArrow(...).mapInArrow(...) produces two MapInArrowExec
operators. The outer
// consumes rows from the inner directly (MapInArrowExec is a row
producer), so there is
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
