feat: hint at native pyarrow UDF path when the feature is disabled

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)

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