fix: correct extended explain operator counts for reuse wrappers and CometSubqueryBroadcast

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,14 +165,31 @@ 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
 inline in the UI. Earlier Spark versions do not have the
 `extendedExplainProviders` extension point, so this provider is not used and
 the config has no effect there.
 
+Not every node in the plan is an eligible operator. The following are excluded
+from both operator counts:
+
+- Transition nodes (`CometColumnarToRow`, `CometNativeColumnarToRow`,
+  `CometSparkRowToColumnar`, `ColumnarToRow`, `RowToColumnar`), which are
+  reported separately as the transition count.
+- Wrappers that do no work of their own: `AdaptiveSparkPlan`, `InputAdapter`,
+  `WholeStageCodegen`, query stages, and `AQEShuffleRead`.
+- The reuse markers `ReusedExchange` and `ReusedSubquery`. The operators they
+  point at are counted once, where the plan they reference is shown, so a
+  reused subtree does not inflate the totals at each reference site.
+
 ### `spark.comet.explain.native.enabled`
 
 When enabled, each executor task logs the DataFusion plan it executes,
@@ -201,12 +218,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 +251,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,16 +19,19 @@
 
 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}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec}
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec}
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
 
-import org.apache.comet.CometExplainInfo.getActualPlan
+import org.apache.comet.CometExplainInfo.{actualPlanForCoverage, getActualPlan}
 import org.apache.comet.annotation.Public
 
 @Public
@@ -43,7 +46,7 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
         // extended information in a tree display.
         val planStats = new CometCoverageStats()
         val outString = new StringBuilder()
-        generateTreeString(getActualPlan(plan), 0, Seq(), 0, outString, planStats)
+        generateTreeString(actualPlanForCoverage(plan), 0, Seq(), 0, outString, planStats)
         s"${outString.toString()}\n$planStats"
       case CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_FALLBACK =>
         // Generates the extended info as a list of fallback reasons
@@ -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)
@@ -105,8 +134,12 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
 
     node match {
       case _: AdaptiveSparkPlanExec | _: InputAdapter | _: QueryStageExec |
-          _: WholeStageCodegenExec | _: ReusedExchangeExec | _: AQEShuffleReadExec =>
-      // ignore
+          _: WholeStageCodegenExec | _: ReusedExchangeExec | _: ReusedSubqueryExec |
+          _: AQEShuffleReadExec =>
+      // Ignore. These nodes either wrap another plan without doing work of their own, or are
+      // pure reuse bookkeeping: the operators they refer to are counted at the site the plan
+      // they reference is shown, so counting them again here would either invent an
+      // un-accelerated Spark operator or double count an accelerated subtree.
       case _: RowToColumnarExec | _: ColumnarToRowExec | _: CometColumnarToRowExec |
           _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec =>
         planStats.transitions += 1
@@ -116,6 +149,8 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
         planStats.sparkOperators += 1
     }
 
+    planStats.recordExpressions(node)
+
     outString.append("   " * indent)
     if (depth > 0) {
       lastChildren.init.foreach { isLast =>
@@ -145,7 +180,7 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
       innerChildrenLocal.init.foreach {
         case c @ (_: TreeNode[_]) =>
           generateTreeString(
-            getActualPlan(c),
+            actualPlanForCoverage(c),
             depth + 2,
             lastChildren :+ node.children.isEmpty :+ false,
             indent,
@@ -154,7 +189,7 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
         case _ =>
       }
       generateTreeString(
-        getActualPlan(innerChildrenLocal.last),
+        actualPlanForCoverage(innerChildrenLocal.last),
         depth + 2,
         lastChildren :+ node.children.isEmpty :+ true,
         indent,
@@ -165,7 +200,7 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
       node.children.init.foreach {
         case c @ (_: TreeNode[_]) =>
           generateTreeString(
-            getActualPlan(c),
+            actualPlanForCoverage(c),
             depth + 1,
             lastChildren :+ false,
             indent,
@@ -176,7 +211,7 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
       node.children.last match {
         case c @ (_: TreeNode[_]) =>
           generateTreeString(
-            getActualPlan(c),
+            actualPlanForCoverage(c),
             depth + 1,
             lastChildren :+ true,
             indent,
@@ -193,13 +228,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."
   }
 }
 
@@ -212,7 +269,7 @@ object CometCoverageStats {
     val stats = new CometCoverageStats()
     val explainInfo = new ExtendedExplainInfo()
     explainInfo.generateTreeString(
-      CometExplainInfo.getActualPlan(plan),
+      CometExplainInfo.actualPlanForCoverage(plan),
       0,
       Seq(),
       0,
@@ -226,21 +283,76 @@ 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[_] = {
+    unwrap(node, unwrapReusedExchange = true)
+  }
+
+  /**
+   * Variant of [[getActualPlan]] used by the coverage traversal in
+   * [[ExtendedExplainInfo.generateTreeString]]. It leaves `ReusedExchangeExec` in place instead
+   * of replacing it with the exchange it points at.
+   *
+   * `ReusedExchangeExec` is a `LeafExecNode`, so keeping the wrapper renders a `ReusedExchange`
+   * leaf and stops the walk there. The reused subtree is then counted once, at the site the
+   * exchange is defined, rather than once per reference. This matches how `ReusedSubqueryExec` is
+   * handled and keeps the rendered tree consistent with the counts reported below it.
+   */
+  def actualPlanForCoverage(node: TreeNode[_]): TreeNode[_] = {
+    unwrap(node, unwrapReusedExchange = false)
+  }
+
+  private def unwrap(node: TreeNode[_], unwrapReusedExchange: Boolean): TreeNode[_] = {
     node match {
-      case p: AdaptiveSparkPlanExec => getActualPlan(p.executedPlan)
-      case p: InputAdapter => getActualPlan(p.child)
-      case p: QueryStageExec => getActualPlan(p.plan)
-      case p: WholeStageCodegenExec => getActualPlan(p.child)
-      case p: ReusedExchangeExec => getActualPlan(p.child)
+      case p: AdaptiveSparkPlanExec => unwrap(p.executedPlan, unwrapReusedExchange)
+      case p: InputAdapter => unwrap(p.child, unwrapReusedExchange)
+      case p: QueryStageExec => unwrap(p.plan, unwrapReusedExchange)
+      case p: WholeStageCodegenExec => unwrap(p.child, unwrapReusedExchange)
+      case p: ReusedExchangeExec if unwrapReusedExchange => unwrap(p.child, unwrapReusedExchange)
       case p => p
     }
-
   }
 
 }