From 544a6bc73df9fffce7fa55b636ba3da6f701e9f4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 11 Jul 2026 02:51:28 -0600 Subject: [PATCH] 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. --- .../rules/EliminateRedundantTransitions.scala | 49 ++++++++++++------- .../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..e886b85fc7 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