chore: Codegen fallback explain for plan

docs/source/user-guide/latest/understanding-comet-plans.md

@@ -84,15 +84,16 @@ incompatibility details.
 
 ## Configs for Inspecting Plans and Fallback
 
-Comet provides four configs for understanding what is happening in a plan.
+Comet provides five configs for understanding what is happening in a plan.
 They serve different purposes and produce output in different places.
 
-| Config                                   | Output destination                 | What you see                                                                                  |
-| ---------------------------------------- | ---------------------------------- | --------------------------------------------------------------------------------------------- |
-| `spark.comet.explainFallback.enabled`    | Driver log (only when fallback)    | A WARN with the list of reasons each query stage could not run in Comet.                      |
-| `spark.comet.logFallbackReasons.enabled` | Driver log                         | One WARN per fallback reason as it is encountered, without surrounding plan context.          |
-| `spark.comet.explain.format`             | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. |
-| `spark.comet.explain.native.enabled`     | Executor logs, per task            | The DataFusion plan with metrics, useful for inspecting Rust execution.                       |
+| Config                                   | Output destination                 | What you see                                                                                                                                                              |
+| ---------------------------------------- | ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `spark.comet.explainFallback.enabled`    | Driver log (only when fallback)    | A WARN with the list of reasons each query stage could not run in Comet.                                                                                                  |
+| `spark.comet.logFallbackReasons.enabled` | Driver log                         | One WARN per fallback reason as it is encountered, without surrounding plan context.                                                                                      |
+| `spark.comet.explainCodegen.enabled`     | `ExtendedExplainInfo` / SQL UI     | A single `[COMET-INFO: JVM codegen dispatcher: <name1>, <name2>, ...]` segment naming each expression on the operator that was routed through the JVM codegen dispatcher. |
+| `spark.comet.explain.format`             | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value.                                                                             |
+| `spark.comet.explain.native.enabled`     | Executor logs, per task            | The DataFusion plan with metrics, useful for inspecting Rust execution.                                                                                                   |
 
 ### `spark.comet.explainFallback.enabled`
 
@@ -108,6 +109,48 @@ when you want to see all reasons, including ones that
 not include the surrounding plan, so it is best for accumulating diagnostics
 across many queries.
 
+### `spark.comet.explainCodegen.enabled`
+
+Disabled by default. When enabled, Comet annotates the surrounding Comet
+operator (`CometProject`, `CometFilter`, etc.) with a single combined
+`[COMET-INFO: JVM codegen dispatcher: <name1>, <name2>, ...]` segment listing
+every expression on that operator that was routed through the JVM codegen
+dispatcher (Spark's own `doGenCode` running inside a Comet kernel), gated by
+`spark.comet.exec.scalaUDF.codegen.enabled=true`. The projection stays
+Comet-native — this is informational only, useful for distinguishing "runs
+natively in DataFusion" from "runs Spark's generated code inside a Comet
+kernel". The annotation only appears for expressions Comet actually routed
+through the dispatcher on this plan — either because no native DataFusion
+implementation exists (`CometCodegenDispatch` serdes such as `hypot`,
+`levenshtein`, `pmod`), or because a native path exists but declined this
+specific input (`CodegenDispatchFallback` mixins). It is not a general "this
+expression ran here" marker.
+
+Example:
+
+```scala
+spark.conf.set("spark.comet.exec.scalaUDF.codegen.enabled", "true")
+spark.conf.set("spark.comet.explainCodegen.enabled", "true")
+
+val df = spark.sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t")
+println(new org.apache.comet.ExtendedExplainInfo()
+  .generateExtendedInfo(df.queryExecution.executedPlan))
+```
+
+Output:
+
+```
+CometNativeColumnarToRow
++- CometProject [COMET-INFO: JVM codegen dispatcher: hypot, levenshtein]
+   +- CometNativeScan parquet spark_catalog.default.t
+
+Comet accelerated 2 out of 2 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet.
+```
+
+Note that the operator is still `CometProject` (Comet-accelerated); only the
+per-expression evaluation for `hypot` and `levenshtein` takes the JVM codegen
+path.
+
 ### `spark.comet.explain.format`
 
 This config is read by `org.apache.comet.ExtendedExplainInfo`, which Spark

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -686,6 +686,15 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_EXPLAIN_CODEGEN_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.explainCodegen.enabled")
+      .category(CATEGORY_EXEC_EXPLAIN)
+      .doc("When enabled, Comet annotates the surrounding Comet operator with a `[COMET-INFO: " +
+        "JVM codegen dispatcher: <names>]` segment listing every expression it routed through " +
+        "the JVM codegen dispatcher. Disabled by default.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_ONHEAP_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.exec.onHeap.enabled")
       .category(CATEGORY_TESTING)

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -383,4 +383,19 @@ object CometSparkSessionExtensions extends Logging {
     node
   }
 
+  /**
+   * Record that `node` (typically an `Expression`) is routing through the JVM codegen dispatcher.
+   * `CometExecRule.rollUpInfoMessages` collects the names across an operator's expression trees
+   * and emits one combined `[COMET-INFO: ...]` segment.
+   */
+  def withCodegenDispatchExpr[T <: TreeNode[_]](node: T, name: String): T = {
+    if (name != null && name.nonEmpty) {
+      val existing = node
+        .getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS)
+        .getOrElse(Set.empty[String])
+      node.setTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS, existing + name)
+    }
+    node
+  }
+
 }

spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala

@@ -224,6 +224,11 @@ object CometExplainInfo {
   val FALLBACK_REASONS = new TreeNodeTag[Set[String]]("CometFallbackReasons")
   val EXTENSION_INFO = new TreeNodeTag[Set[String]]("CometExtensionInfo")
 
+  // Expression names the serde routed through the JVM codegen dispatcher. Rolled up per
+  // operator by `CometExecRule.rollUpInfoMessages` into one combined `[COMET-INFO: ...]`.
+  val CODEGEN_DISPATCH_EXPRS =
+    new TreeNodeTag[Set[String]]("CometCodegenDispatchExprs")
+
   def getActualPlan(node: TreeNode[_]): TreeNode[_] = {
     node match {
       case p: AdaptiveSparkPlanExec => getActualPlan(p.executedPlan)

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -729,14 +729,25 @@ case class CometExecRule(session: SparkSession)
    * Lift informational (non-fallback) messages tagged on an operator and its expressions onto the
    * converted Comet plan node so they appear in verbose extended explain output. Expression-level
    * hints would otherwise be invisible because explain only traverses plan nodes, not
-   * expressions.
+   * expressions. `CODEGEN_DISPATCH_EXPRS` names across the tree are aggregated into one combined
+   * info line.
    */
   private def rollUpInfoMessages(op: SparkPlan, exec: SparkPlan): Unit = {
-    val fromOp = op.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String])
-    val fromExprs = op.expressions
-      .flatMap(_.collect { case e: Expression => e })
-      .flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String]))
-    (fromOp ++ fromExprs).foreach(msg => withInfo(exec, msg))
+    val allExprs = op.expressions.flatMap(_.collect { case e: Expression => e })
+
+    val infos =
+      op.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String]) ++
+        allExprs.flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO)).flatten
+    infos.foreach(msg => withInfo(exec, msg))
+
+    val routedNames = allExprs
+      .flatMap(_.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS))
+      .flatten
+      .distinct
+      .sorted
+    if (routedNames.nonEmpty) {
+      withInfo(exec, s"JVM codegen dispatcher: ${routedNames.mkString(", ")}")
+    }
   }
 
   private def isOperatorEnabled(

spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala

@@ -19,12 +19,14 @@
 
 package org.apache.comet.serde
 
+import java.util.Locale
+
 import org.apache.spark.SparkEnv
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, RuntimeReplaceable, ScalaUDF}
 import org.apache.spark.sql.types.BinaryType
 
 import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
+import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, withFallbackReason}
 import org.apache.comet.codegen.CometBatchKernelCodegen
 import org.apache.comet.serde.ExprOuterClass.Expr
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType}
@@ -55,6 +57,19 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
   override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] =
     emitJvmCodegenDispatch(expr, inputs, binding)
 
+  /**
+   * Canonical name for annotation. `BinaryMathExpression` (Hypot, Pow, ...) overrides
+   * `prettyName` to raw uppercase; `ScalaUDF.prettyName` collapses to `"scalaudf"` for every user
+   * UDF. Prefer `ScalaUDF.udfName` when set, then lowercase to normalize.
+   */
+  private def displayName(expr: Expression): String = {
+    val raw = expr match {
+      case s: ScalaUDF => s.udfName.getOrElse(s.prettyName)
+      case other => other.prettyName
+    }
+    raw.toLowerCase(Locale.ROOT)
+  }
+
   /**
    * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through
    * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the
@@ -70,11 +85,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
       expr: Expression,
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] = {
+    val exprName = displayName(expr)
     if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) {
       withFallbackReason(
         expr,
-        s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has no native " +
-          "path so the plan falls back to Spark")
+        s"$exprName: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has " +
+          "no native path so the plan falls back to Spark")
       return None
     }
 
@@ -97,7 +113,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
     // at execute.
     CometBatchKernelCodegen.canHandle(boundExpr) match {
       case Some(reason) =>
-        withFallbackReason(expr, reason)
+        withFallbackReason(expr, s"$exprName: $reason")
         return None
       case None =>
     }
@@ -110,12 +126,26 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
     val buffer = serializer.serialize(boundExpr)
     val bytes = new Array[Byte](buffer.remaining())
     buffer.get(bytes)
-    val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding)
-      .getOrElse(return None)
+    val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding).getOrElse {
+      withFallbackReason(
+        expr,
+        s"$exprName: codegen dispatch: could not serialize closure-serialized bound " +
+          "expression payload")
+      return None
+    }
 
-    val dataArgs =
-      attrs.map(a => exprToProtoInternal(a, inputs, binding).getOrElse(return None))
-    val returnTypeProto = serializeDataType(expr.dataType).getOrElse(return None)
+    val dataArgs = attrs.map { a =>
+      exprToProtoInternal(a, inputs, binding).getOrElse {
+        withFallbackReason(expr, s"$exprName: codegen dispatch: could not serialize data arg $a")
+        return None
+      }
+    }
+    val returnTypeProto = serializeDataType(expr.dataType).getOrElse {
+      withFallbackReason(
+        expr,
+        s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}")
+      return None
+    }
 
     val udfBuilder = ExprOuterClass.JvmScalarUdf
       .newBuilder()
@@ -125,6 +155,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
     udfBuilder
       .setReturnType(returnTypeProto)
       .setReturnNullable(expr.nullable)
+    // Opt-in dispatch annotation for extended explain. Rolled up per operator by
+    // `CometExecRule.rollUpInfoMessages` into a single `[COMET-INFO: JVM codegen dispatcher:
+    // ...]` line. Informational only - does not trigger fallback.
+    if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) {
+      withCodegenDispatchExpr(expr, exprName)
+    }
     Some(
       ExprOuterClass.Expr
         .newBuilder()

spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala

@@ -131,6 +131,72 @@ class CometCodegenSuite
     }
   }
 
+  test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") {
+    // With the opt-in flag on, `hypot` and `levenshtein` (both `CometCodegenDispatch`) roll up
+    // into one `[COMET-INFO: JVM codegen dispatcher: hypot, levenshtein]` line on the
+    // `CometProject`. With the flag off (default), no such line appears.
+    withTable("t") {
+      sql("CREATE TABLE t (a DOUBLE, b DOUBLE, s1 STRING, s2 STRING) USING parquet")
+      sql("INSERT INTO t VALUES (3.0, 4.0, 'kitten', 'sitting')")
+
+      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) {
+        val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t")
+        checkSparkAnswerAndOperator(df)
+        val explain =
+          new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan)
+        assert(
+          explain.contains("[COMET-INFO:"),
+          s"expected a [COMET-INFO: segment, got:\n$explain")
+        // Names appear alphabetically via `.distinct.sorted` in rollUpInfoMessages.
+        assert(
+          explain.contains("JVM codegen dispatcher: hypot, levenshtein"),
+          s"expected combined codegen-dispatch info, got:\n$explain")
+      }
+
+      withSQLConf(
+        CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "false",
+        CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true",
+        CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key ->
+          CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) {
+        val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t")
+        checkSparkAnswerAndOperator(df)
+        val explain =
+          new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan)
+        assert(
+          !explain.contains("JVM codegen dispatcher"),
+          s"expected NO codegen-dispatch info with the flag off, got:\n$explain")
+      }
+    }
+  }
+
+  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.
+    withTable("t") {
+      sql("CREATE TABLE t (a DOUBLE, b DOUBLE) USING parquet")
+      sql("INSERT INTO t VALUES (3.0, 4.0)")
+      withSQLConf(
+        CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false",
+        CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key ->
+          CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) {
+        val df = sql("SELECT hypot(a, b) FROM t")
+        checkSparkAnswer(df)
+        val explain =
+          new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan)
+        assert(
+          explain.contains("hypot:") &&
+            explain.contains(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false"),
+          s"expected 'hypot:' prefix and disabled-flag reason, got:\n$explain")
+      }
+    }
+  }
+
   test("dispatcher caches the compiled kernel across batches of one query") {
     // Within a single query, the dispatcher compiles a kernel for the (expression, schema) pair
     // once and reuses it across every subsequent batch of the same shape. Force multiple batches