feat: report native vs codegen-dispatch expression coverage in extended explain

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

@@ -144,7 +144,7 @@ CometColumnarToRow
 +- 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.
+Comet accelerated 2 out of 2 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 0 native, 2 codegen dispatch.
 ```
 
 Note that the operator is still `CometProject` (Comet-accelerated); only the
@@ -165,7 +165,12 @@ The Spark SQL UI then shows an additional section under the detailed plan.
 The format is controlled by `spark.comet.explain.format`:
 
 - `verbose` (default): the full plan annotated with fallback reasons, plus a
-  summary of how much of the plan is accelerated.
+  summary of how much of the plan is accelerated. The summary reports operator
+  coverage, the number of Spark/Comet transitions, and how many distinct
+  expressions Comet accelerated - split into those lowered to native DataFusion
+  expressions and those routed through the JVM codegen dispatcher. Expression
+  names are counted once per plan, and structural nodes (attribute references,
+  literals, aliases) are excluded.
 - `fallback`: a list of fallback reasons only.
 
 This is the most convenient option on Spark 4.0 because the output is shown
@@ -201,12 +206,22 @@ val info = new ExtendedExplainInfo()
 // Sorted, deduplicated list of fallback reasons across the whole plan.
 val reasons: Seq[String] = info.getFallbackReasons(plan)
 
+// Sorted, deduplicated names of the expressions Comet accelerated, split by how
+// they run. Structural nodes (attribute references, literals, aliases) are not
+// reported.
+val native: Seq[String] = info.getNativeExpressions(plan)
+val dispatched: Seq[String] = info.getCodegenDispatchExpressions(plan)
+
 // Formatted string. Honors spark.comet.explain.format:
 //   - "verbose"  -> the full plan annotated with per-node fallback reasons
 //   - "fallback" -> just the list of reasons
 val formatted: String = info.generateExtendedInfo(plan)
 ```
 
+A name can appear in both lists: the same function may be lowered natively for
+one set of arguments and routed through the dispatcher for another elsewhere in
+the same plan.
+
 Example:
 
 ```
@@ -224,7 +239,7 @@ Project [COMET: from_unixtime(eventTime#5L, yyyy-MM-dd HH:mm:ss, Some(GMT)) is n
 +- CometColumnarToRow
    +- CometNativeScan parquet
 
-Comet accelerated 1 out of 2 eligible operators (50%). Final plan contains 1 transitions between Spark and Comet.
+Comet accelerated 1 out of 2 eligible operators (50%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 0 native, 0 codegen dispatch.
 ```
 
 ## Comet Operator Reference

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

@@ -26,7 +26,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.network.util.ByteUnit
 import org.apache.spark.sql.{SparkSession, SparkSessionExtensions}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreeNode
+import org.apache.spark.sql.catalyst.trees.{TreeNode, TreeNodeTag}
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.internal.SQLConf
@@ -375,12 +375,7 @@ object CometSparkSessionExtensions extends Logging {
    * implementation gated behind a config.
    */
   def withInfo[T <: TreeNode[_]](node: T, message: String): T = {
-    if (message != null && message.nonEmpty) {
-      val existing =
-        node.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String])
-      node.setTagValue(CometExplainInfo.EXTENSION_INFO, existing + message)
-    }
-    node
+    appendTagValue(node, CometExplainInfo.EXTENSION_INFO, message)
   }
 
   /**
@@ -389,11 +384,41 @@ object CometSparkSessionExtensions extends Logging {
    * 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)
+    appendTagValue(node, CometExplainInfo.CODEGEN_DISPATCH_EXPRS, name)
+  }
+
+  /**
+   * Record that `node` (typically an `Expression`) was lowered to a native DataFusion expression.
+   * The native counterpart of [[withCodegenDispatchExpr]]: `CometExecRule.rollUpInfoMessages`
+   * collects the names across an operator's expression trees onto the converted Comet plan node,
+   * where extended explain reads them for expression coverage stats.
+   */
+  def withNativeExpr[T <: TreeNode[_]](node: T, name: String): T = {
+    appendTagValue(node, CometExplainInfo.NATIVE_EXPRS, name)
+  }
+
+  /**
+   * Add `value` to a `Set`-valued `TreeNodeTag`, accumulating rather than overwriting. Null and
+   * empty values are dropped so callers do not have to guard. Shared by [[withInfo]] and the
+   * expression coverage tags.
+   */
+  private def appendTagValue[T <: TreeNode[_]](
+      node: T,
+      tag: TreeNodeTag[Set[String]],
+      value: String): T = {
+    if (value != null && value.nonEmpty) {
+      appendTagValues(node, tag, Set(value))
+    }
+    node
+  }
+
+  /** Bulk form of [[appendTagValue]], for lifting a whole name set onto another node. */
+  private[comet] def appendTagValues[T <: TreeNode[_]](
+      node: T,
+      tag: TreeNodeTag[Set[String]],
+      values: Set[String]): T = {
+    if (values.nonEmpty) {
+      node.setTagValue(tag, node.getTagValue(tag).getOrElse(Set.empty[String]) ++ values)
     }
     node
   }

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

@@ -19,9 +19,12 @@
 
 package org.apache.comet
 
+import java.util.Locale
+
 import scala.collection.mutable
 
 import org.apache.spark.sql.ExtendedExplainGenerator
+import org.apache.spark.sql.catalyst.expressions.{Expression, ScalaUDF}
 import org.apache.spark.sql.catalyst.trees.{TreeNode, TreeNodeTag}
 import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometPlan, CometSparkToColumnarExec}
 import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter, RowToColumnarExec, SparkPlan, WholeStageCodegenExec}
@@ -55,6 +58,32 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
     fallbackReasons(plan).toSeq.sorted
   }
 
+  /**
+   * Names of the expressions in `plan` that Comet lowered to native DataFusion expressions,
+   * sorted alphabetically. Names are the expression's `prettyName` (the UDF name for a
+   * `ScalaUDF`) lowercased.
+   *
+   * Structural nodes - attribute references, literals, aliases, bound references - are not
+   * reported: they carry no computation and would swamp the interesting names.
+   */
+  def getNativeExpressions(plan: SparkPlan): Seq[String] = {
+    exprCoverage(plan, CometExplainInfo.NATIVE_EXPRS)
+  }
+
+  /**
+   * Names of the expressions in `plan` that Comet kept inside the native pipeline by routing them
+   * through the JVM codegen dispatcher (Spark's own `doGenCode` compiled into a batch kernel)
+   * rather than lowering them to a native DataFusion expression, sorted alphabetically. See
+   * [[getNativeExpressions]] for how names are derived.
+   */
+  def getCodegenDispatchExpressions(plan: SparkPlan): Seq[String] = {
+    exprCoverage(plan, CometExplainInfo.CODEGEN_DISPATCH_EXPRS)
+  }
+
+  private def exprCoverage(plan: SparkPlan, tag: TreeNodeTag[Set[String]]): Seq[String] = {
+    CometExplainInfo.collectTagValues(sortup(plan).toSeq, tag).toSeq.sorted
+  }
+
   private[comet] def fallbackReasons(node: TreeNode[_]): Set[String] = {
     var info = mutable.Seq[String]()
     val sorted = sortup(node)
@@ -116,6 +145,8 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
         planStats.sparkOperators += 1
     }
 
+    planStats.recordExpressions(node)
+
     outString.append("   " * indent)
     if (depth > 0) {
       lastChildren.init.foreach { isLast =>
@@ -193,13 +224,35 @@ class CometCoverageStats {
   var cometOperators: Int = 0
   var transitions: Int = 0
 
+  /** Distinct names of expressions lowered to native DataFusion expressions. */
+  val nativeExpressions: mutable.Set[String] = mutable.HashSet.empty
+
+  /** Distinct names of expressions routed through the JVM codegen dispatcher. */
+  val codegenDispatchExpressions: mutable.Set[String] = mutable.HashSet.empty
+
+  /**
+   * Accumulate the expression coverage that `CometExecRule.rollUpInfoMessages` rolled up onto a
+   * converted Comet plan node.
+   */
+  private[comet] def recordExpressions(node: TreeNode[_]): Unit = {
+    node.getTagValue(CometExplainInfo.NATIVE_EXPRS).foreach(nativeExpressions ++= _)
+    node
+      .getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS)
+      .foreach(codegenDispatchExpressions ++= _)
+  }
+
   override def toString(): String = {
     val eligible = sparkOperators + cometOperators
     val converted =
       if (eligible == 0) 0.0 else cometOperators.toDouble / eligible * 100.0
+    // Deliberately no combined expression total: the counts are of distinct names, and the same
+    // function can be lowered natively for one set of arguments and dispatched for another, so a
+    // name can be in both. A total would not be the sum and would read as an arithmetic error.
     s"Comet accelerated $cometOperators out of $eligible " +
       s"eligible operators (${converted.toInt}%). " +
-      s"Final plan contains $transitions transitions between Spark and Comet."
+      s"Final plan contains $transitions transitions between Spark and Comet. " +
+      s"Accelerated expressions: ${nativeExpressions.size} native, " +
+      s"${codegenDispatchExpressions.size} codegen dispatch."
   }
 }
 
@@ -226,11 +279,49 @@ 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: ...]`.
+  // Expression names the serde routed through the JVM codegen dispatcher. Set on each such
+  // expression, then rolled up per operator by `CometExecRule.rollUpInfoMessages` onto the
+  // converted Comet plan node (where extended explain reads it for coverage stats, and where it
+  // becomes one combined `[COMET-INFO: ...]` when `spark.comet.explain.codegen.enabled` is set).
+  // Because the tag accumulates names from descendants as well, it answers "was anything under
+  // here dispatched", not "was this node itself dispatched" - use `DISPATCHED_SELF` for that.
   val CODEGEN_DISPATCH_EXPRS =
     new TreeNodeTag[Set[String]]("CometCodegenDispatchExprs")
 
+  // Marks the one expression that `CometScalaUDF.emitJvmCodegenDispatch` converted, as opposed to
+  // an ancestor that merely carries a descendant's name in `CODEGEN_DISPATCH_EXPRS`. Never rolled
+  // up or lifted, so classifying an expression as native/dispatched does not depend on the order
+  // in which the planner happens to visit it.
+  val DISPATCHED_SELF = new TreeNodeTag[Unit]("CometDispatchedSelf")
+
+  // Expression names the serde lowered to native DataFusion expressions. The native counterpart
+  // of `CODEGEN_DISPATCH_EXPRS`, rolled up the same way, but never rendered in the tree display:
+  // it would repeat what the operator names already say.
+  val NATIVE_EXPRS = new TreeNodeTag[Set[String]]("CometNativeExprs")
+
+  /**
+   * Union of a `Set`-valued tag over `nodes`. Used to roll expression coverage names up onto an
+   * operator and to gather them back off a plan.
+   */
+  def collectTagValues(nodes: Seq[TreeNode[_]], tag: TreeNodeTag[Set[String]]): Set[String] = {
+    nodes.flatMap(_.getTagValue(tag).getOrElse(Set.empty[String])).toSet
+  }
+
+  /**
+   * Name used to report `expr` in explain output.
+   *
+   * `BinaryMathExpression` (Hypot, Pow, ...) overrides `prettyName` to raw uppercase; `ScalaUDF`
+   * collapses to `"scalaudf"` for every user UDF. Prefer `ScalaUDF.udfName` when set, then
+   * lowercase to normalize.
+   */
+  def exprDisplayName(expr: Expression): String = {
+    val raw = expr match {
+      case s: ScalaUDF => s.udfName.getOrElse(s.prettyName)
+      case other => other.prettyName
+    }
+    raw.toLowerCase(Locale.ROOT)
+  }
+
   def getActualPlan(node: TreeNode[_]): TreeNode[_] = {
     node match {
       case p: AdaptiveSparkPlanExec => getActualPlan(p.executedPlan)

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

@@ -745,24 +745,28 @@ 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. `CODEGEN_DISPATCH_EXPRS` names across the tree are aggregated into one combined
-   * info line.
+   * expressions. `NATIVE_EXPRS` and `CODEGEN_DISPATCH_EXPRS` names across the tree are lifted the
+   * same way so that extended explain can report expression coverage; the dispatched ones
+   * additionally become one combined info line when `spark.comet.explain.codegen.enabled` is set.
    */
   private def rollUpInfoMessages(op: SparkPlan, exec: SparkPlan): Unit = {
     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
+        CometExplainInfo.collectTagValues(allExprs, CometExplainInfo.EXTENSION_INFO)
     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(", ")}")
+    appendTagValues(
+      exec,
+      CometExplainInfo.NATIVE_EXPRS,
+      CometExplainInfo.collectTagValues(allExprs, CometExplainInfo.NATIVE_EXPRS))
+
+    val routedNames =
+      CometExplainInfo.collectTagValues(allExprs, CometExplainInfo.CODEGEN_DISPATCH_EXPRS)
+    appendTagValues(exec, CometExplainInfo.CODEGEN_DISPATCH_EXPRS, routedNames)
+    if (routedNames.nonEmpty && CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) {
+      withInfo(exec, s"JVM codegen dispatcher: ${routedNames.toSeq.sorted.mkString(", ")}")
     }
   }
 

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

@@ -19,13 +19,12 @@
 
 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.CometExplainInfo
 import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, withFallbackReason}
 import org.apache.comet.codegen.CometBatchKernelCodegen
 import org.apache.comet.serde.ExprOuterClass.Expr
@@ -57,19 +56,6 @@ 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
@@ -85,7 +71,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] {
       expr: Expression,
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] = {
-    val exprName = displayName(expr)
+    val exprName = CometExplainInfo.exprDisplayName(expr)
     if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) {
       withFallbackReason(
         expr,
@@ -155,12 +141,14 @@ 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)
-    }
+    // Dispatch annotation for extended explain. Rolled up per operator by
+    // `CometExecRule.rollUpInfoMessages`, which feeds the expression coverage stats and, when
+    // `spark.comet.explain.codegen.enabled` is set, a single `[COMET-INFO: JVM codegen dispatcher:
+    // ...]` line. Informational only - does not trigger fallback. The marker records that this
+    // node itself was dispatched, which the name set alone cannot say once ancestors accumulate
+    // their descendants' names.
+    expr.setTagValue(CometExplainInfo.DISPATCHED_SELF, ())
+    withCodegenDispatchExpr(expr, exprName)
     Some(
       ExprOuterClass.Expr
         .newBuilder()