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 `[COMET-INFO: <expr>: routes through JVM codegen dispatcher]` annotation per routed expression on the surrounding Comet operator. |
+| `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,44 @@ 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, every Spark expression that Comet routes
+through the JVM codegen dispatcher (Spark's own `doGenCode` running inside a
+Comet kernel — the path used for expressions with no native DataFusion
+implementation, gated by `spark.comet.exec.scalaUDF.codegen.enabled=true`) is
+annotated with `<expr_name>: routes through JVM codegen dispatcher`. The
+annotation rolls up onto the surrounding `CometProject` / `CometFilter` /
+etc. and is rendered inside the `[COMET-INFO: ...]` segment of extended
+explain output (verbose format). 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".
+
+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: hypot: routes through JVM codegen dispatcher, levenshtein: routes through JVM codegen dispatcher]
+   +- 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,17 @@ 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 this setting is enabled, Comet will annotate expressions that route through the " +
+          "JVM codegen dispatcher (Spark's `doGenCode` running inside a Comet kernel) with a " +
+          "`[COMET-INFO: ...]` message in extended explain output. Disabled by default to keep " +
+          "the plan display quiet unless the user is investigating dispatch coverage.")
+      .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/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.{withFallbackReason, withInfo}
 import org.apache.comet.codegen.CometBatchKernelCodegen
 import org.apache.comet.serde.ExprOuterClass.Expr
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType}
@@ -70,11 +72,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
       expr: Expression,
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] = {
+    val exprName = expr.prettyName
     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 +100,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 +113,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 +142,17 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
     udfBuilder
       .setReturnType(returnTypeProto)
       .setReturnNullable(expr.nullable)
+    // Surface the dispatch route in extended explain output when the user opts in via
+    // `spark.comet.explainCodegen.enabled=true`. `rollUpInfoMessages` in `CometExecRule` walks
+    // `op.expressions` via `TreeNode.collect` and copies each expression's `EXTENSION_INFO` tag
+    // onto the resulting Comet plan node, where `ExtendedExplainInfo` renders it as
+    // `[COMET-INFO: ...]`. Informational only: `withInfo` does not trigger fallback, so the
+    // projection remains a Comet operator regardless of this setting.
+    if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) {
+      withInfo(
+        expr,
+        s"${exprName.toLowerCase(Locale.ROOT)}: routes through JVM codegen dispatcher")
+    }
     Some(
       ExprOuterClass.Expr
         .newBuilder()

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

@@ -131,6 +131,42 @@ class CometCodegenSuite
     }
   }
 
+  test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") {
+    // Two built-in expressions in the same SELECT that both route through the JVM codegen
+    // dispatcher (hypot -> CometHypot, levenshtein -> CometLevenshtein, both extend
+    // CometCodegenDispatch). With codegen dispatch enabled AND
+    // `spark.comet.explainCodegen.enabled=true`, `emitJvmCodegenDispatch` tags each expression
+    // with a withInfo message on its EXTENSION_INFO tag. `CometExecRule.rollUpInfoMessages`
+    // walks the projection's expression trees and copies those tags onto the CometProjectExec
+    // node, where `ExtendedExplainInfo` renders them as `[COMET-INFO: ...]`. This test verifies
+    // both named expressions surface in the info block while the projection stays Comet-native.
+    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 in explain output, got:\n$explain")
+        assert(
+          explain.contains("hypot: routes through JVM codegen dispatcher"),
+          s"expected 'hypot' codegen-dispatch info in explain output, got:\n$explain")
+        assert(
+          explain.contains("levenshtein: routes through JVM codegen dispatcher"),
+          s"expected 'levenshtein' codegen-dispatch info in explain output, 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