refactor: replace hand-coded rollup of expression fallback reasons onto operators

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

@@ -668,6 +668,20 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_STRICT_FALLBACK_REASONS: ConfigEntry[Boolean] =
+    conf("spark.comet.explain.fallback.strict.enabled")
+      .category(CATEGORY_TESTING)
+      .doc(
+        "Test-only. When enabled, Comet throws if it declines to convert an operator that it " +
+          "could otherwise have converted (all children are already native) without recording a " +
+          "fallback reason on the operator or on any of its expressions. Without this check, a " +
+          "serde that returns `None` and forgets to state a reason silently produces a generic " +
+          "'<operator> is not supported' message instead of a visible failure. Enabled for all " +
+          "Comet test suites via `CometTestBase`.")
+      .internal()
+      .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

@@ -274,41 +274,41 @@ object CometSparkSessionExtensions extends Logging {
    * Record a fallback reason on a `TreeNode` (a Spark operator or expression) explaining why
    * Comet cannot accelerate it. Reasons recorded here are surfaced in extended explain output
    * (see `ExtendedExplainInfo`) and, when `COMET_EXPLAIN_FALLBACK_LOG_ENABLED` is enabled, logged
-   * as warnings. The reasons are also rolled up from child nodes so that the operator that
-   * remains in the Spark plan carries the reasons from its converted-away subtree.
+   * as warnings.
    *
    * Call this in any code path where Comet decides not to convert a given node - serde `convert`
    * methods returning `None`, unsupported data types, disabled configs, etc. Do not use this for
    * informational messages that are not fallback reasons: anything tagged here is treated by the
    * rules as a signal that the node falls back to Spark.
    *
+   * Tag only the node that actually failed, and state a real reason. There is deliberately no way
+   * to copy reasons from child nodes onto a parent: extended explain only walks plan nodes, so an
+   * expression-level reason is lifted onto the enclosing operator centrally by
+   * `CometExecRule.rollUpFallbackReasons` when that operator is left in the Spark plan. See
+   * https://github.com/apache/datafusion-comet/issues/5230.
+   *
    * @param node
    *   The Spark operator or expression that is falling back to Spark.
    * @param info
-   *   The fallback reason. Optional, may be null or empty - pass empty only when the call is used
-   *   purely to roll up reasons from `exprs`.
-   * @param exprs
-   *   Child nodes whose own fallback reasons should be rolled up into `node`. Pass the
-   *   sub-expressions or child operators whose failure caused `node` to fall back.
+   *   The fallback reason. Newline-delimited to record more than one reason.
    * @tparam T
    *   The type of the TreeNode. Typically `SparkPlan`, `AggregateExpression`, or `Expression`.
    * @return
-   *   `node` with fallback reasons attached (as a side effect on its tag map).
+   *   `node` with the fallback reason attached (as a side effect on its tag map).
    */
-  def withFallbackReason[T <: TreeNode[_]](node: T, info: String, exprs: T*): T = {
+  def withFallbackReason[T <: TreeNode[_]](node: T, info: String): T = {
     // support existing approach of passing in multiple infos in a newline-delimited string
     val infoSet = if (info == null || info.isEmpty) {
       Set.empty[String]
     } else {
       info.split("\n").toSet
     }
-    withFallbackReasons(node, infoSet, exprs: _*)
+    withFallbackReasons(node, infoSet)
   }
 
   /**
-   * Record one or more fallback reasons on a `TreeNode` and roll up reasons from any child nodes.
-   * This is the set-valued form of [[withFallbackReason]]; see that overload for the full
-   * contract.
+   * Record one or more fallback reasons on a `TreeNode`. This is the set-valued form of
+   * [[withFallbackReason]]; see that overload for the full contract.
    *
    * Reasons are accumulated (never overwritten) on the node's `FALLBACK_REASONS` tag and are
    * surfaced in extended explain output. When `COMET_EXPLAIN_FALLBACK_LOG_ENABLED` is enabled,
@@ -317,50 +317,32 @@ object CometSparkSessionExtensions extends Logging {
    * @param node
    *   The Spark operator or expression that is falling back to Spark.
    * @param info
-   *   The fallback reasons for this node. May be empty when the call is used purely to roll up
-   *   child reasons.
-   * @param exprs
-   *   Child nodes whose own fallback reasons should be rolled up into `node`.
+   *   The fallback reasons for this node.
    * @tparam T
    *   The type of the TreeNode. Typically `SparkPlan`, `AggregateExpression`, or `Expression`.
    * @return
    *   `node` with fallback reasons attached (as a side effect on its tag map).
    */
-  def withFallbackReasons[T <: TreeNode[_]](node: T, info: Set[String], exprs: T*): T = {
+  def withFallbackReasons[T <: TreeNode[_]](node: T, info: Set[String]): T = {
     if (CometConf.COMET_EXPLAIN_FALLBACK_LOG_ENABLED.get()) {
       for (reason <- info) {
         logWarning(s"Comet cannot accelerate ${node.getClass.getSimpleName} because: $reason")
       }
     }
-    val existingNodeInfos = node.getTagValue(CometExplainInfo.FALLBACK_REASONS)
-    val newNodeInfo = (existingNodeInfos ++ exprs
-      .flatMap(_.getTagValue(CometExplainInfo.FALLBACK_REASONS))).flatten.toSet
-    node.setTagValue(CometExplainInfo.FALLBACK_REASONS, newNodeInfo ++ info)
+    val existingNodeInfos =
+      node.getTagValue(CometExplainInfo.FALLBACK_REASONS).getOrElse(Set.empty[String])
+    node.setTagValue(CometExplainInfo.FALLBACK_REASONS, existingNodeInfos ++ info)
     node
   }
 
-  /**
-   * Roll up fallback reasons from `exprs` onto `node` without adding a new reason of its own. Use
-   * this when a parent operator is itself falling back and wants to preserve the reasons recorded
-   * on its child expressions/operators so they appear together in explain output.
-   *
-   * @param node
-   *   The parent operator or expression falling back to Spark.
-   * @param exprs
-   *   Child nodes whose fallback reasons should be aggregated onto `node`.
-   * @tparam T
-   *   The type of the TreeNode. Typically `SparkPlan`, `AggregateExpression`, or `Expression`.
-   * @return
-   *   `node` with the rolled-up reasons attached (as a side effect on its tag map).
-   */
-  def withFallbackReason[T <: TreeNode[_]](node: T, exprs: T*): T = {
-    withFallbackReasons(node, Set.empty, exprs: _*)
-  }
-
   /**
    * True if any fallback reason has been recorded on `node` (via [[withFallbackReason]] /
    * [[withFallbackReasons]]). Callers that need to short-circuit when a prior rule pass has
    * already decided a node falls back can use this as the sticky signal.
+   *
+   * This deliberately reads only the node's own tag. It is a planning control signal, not explain
+   * output, so it must not observe reasons that merely exist somewhere in the node's expression
+   * trees - see `CometExecRule.rollUpFallbackReasons`.
    */
   def hasFallbackReason(node: TreeNode[_]): Boolean = {
     node.getTagValue(CometExplainInfo.FALLBACK_REASONS).exists(_.nonEmpty)

spark/src/main/scala/org/apache/comet/expressions/CometCast.scala

@@ -118,7 +118,6 @@ object CometCast
           if (childExpr.isDefined) {
             castToProto(cast, cast.timeZoneId, cast.dataType, childExpr.get, cometEvalMode)
           } else {
-            withFallbackReason(cast, cast.child)
             None
           }
         }

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

@@ -382,7 +382,11 @@ case class CometExecRule(session: SparkSession)
             // when COMET_EXPLAIN_FALLBACK_LOG_ENABLED=true) even when the write is fully native.
             op
           case _ =>
-            // The operator was not converted to a Comet plan. Possible reasons for this happening:
+            // The operator was not converted to a Comet plan and no serde handler claimed it, so
+            // Comet simply has no support for it. (Operators that do have a handler are reported
+            // by `reportUnexplainedFallback` inside `convertToComet`, which is also where the
+            // strict check lives - it would be wrong to demand a specific reason here, because
+            // nothing ever attempted this operator.) Possible reasons for reaching this point:
             // 1. Comet does not support this operator.
             // 2. The operator could not be supported based on query context and current
             //    configs. In this case, it should have already been tagged with fallback
@@ -698,6 +702,23 @@ case class CometExecRule(session: SparkSession)
 
   /** Convert a Spark plan to a Comet plan using the specified serde handler */
   private def convertToComet(op: SparkPlan, handler: CometOperatorSerde[_]): Option[SparkPlan] = {
+    val converted = tryConvertToComet(op, handler)
+    if (converted.isEmpty) {
+      // Comet looked at this operator and declined it, so it stays in the Spark plan. Lift any
+      // reasons recorded on its expressions onto the operator itself - see
+      // `rollUpFallbackReasons` for why this is needed - and then make sure something was
+      // recorded. The order is required, not incidental: `reportUnexplainedFallback` inspects only
+      // the operator's own tag, so a reason still sitting on an expression would look like no
+      // reason at all and trip the strict check.
+      rollUpFallbackReasons(op)
+      reportUnexplainedFallback(op)
+    }
+    converted
+  }
+
+  private def tryConvertToComet(
+      op: SparkPlan,
+      handler: CometOperatorSerde[_]): Option[SparkPlan] = {
     val serde = handler.asInstanceOf[CometOperatorSerde[SparkPlan]]
     if (isOperatorEnabled(serde, op)) {
       // For operators that require native children (like writes), check if all data-producing
@@ -741,6 +762,69 @@ case class CometExecRule(session: SparkSession)
     None
   }
 
+  /**
+   * Lift fallback reasons recorded on `op`'s expression trees onto `op` itself.
+   *
+   * Extended explain output only walks plan nodes (`ExtendedExplainInfo.sortup` follows
+   * `children` / `innerChildren`, never `expressions`), so a reason tagged on an expression is
+   * invisible unless something lifts it onto the enclosing operator. This mirrors what
+   * [[rollUpInfoMessages]] already does for the informational tags, and replaces the roll-up that
+   * used to be hand-written at every serde call site (see
+   * https://github.com/apache/datafusion-comet/issues/5230).
+   *
+   * Only child *expressions* are collected, not child operators: reasons on a child operator are
+   * already reachable by the explain traversal via `children`.
+   *
+   * Called only when `op` was left in the Spark plan, which scopes the roll-up to the operator
+   * that actually failed conversion. That matters because some expression instances
+   * (`AttributeReference`s, DPP subquery expressions) are shared across operators, so an unscoped
+   * roll-up could surface one expression's reason under several unrelated operators.
+   *
+   * [[reportUnexplainedFallback]] relies on this having run first; the two must not be separated.
+   */
+  private def rollUpFallbackReasons(op: SparkPlan): Unit = {
+    val reasons = op.expressions
+      .flatMap(_.collect { case e: Expression => e })
+      .flatMap(_.getTagValue(CometExplainInfo.FALLBACK_REASONS))
+      .flatten
+      .toSet
+    if (reasons.nonEmpty) {
+      withFallbackReasons(op, reasons)
+    }
+  }
+
+  /**
+   * Handle an operator that Comet declined without stating why.
+   *
+   * When every child is already native, Comet had a real opportunity to convert `op`, so the
+   * absence of any reason - on `op` or anywhere in its expression trees - means a serde returned
+   * `None` and forgot to record one. Under `COMET_STRICT_FALLBACK_REASONS` (enabled for Comet's
+   * own test suites) that is a hard failure; otherwise fall back to a generic message so users
+   * still see something. The generic message is what used to mask this whole class of bug, which
+   * is why the strict check exists.
+   *
+   * Must run *after* [[rollUpFallbackReasons]] for the same operator. The check reads only `op`'s
+   * own tag, because `hasFallbackReason` deliberately does not traverse expressions (it is a
+   * planning control signal, not explain output), so an expression-level reason that has not been
+   * lifted yet would be mistaken for no reason at all. [[convertToComet]] is the only production
+   * caller and keeps the two calls together.
+   *
+   * Package-visible so `CometExecRuleSuite` can drive the strict failure directly: no serde in
+   * the tree reaches this state, which is exactly what the check enforces, so the only way to
+   * test it is to construct the shape by hand.
+   */
+  private[comet] def reportUnexplainedFallback(op: SparkPlan): Unit = {
+    if (op.children.forall(_.isInstanceOf[CometNativeExec]) && !hasFallbackReason(op)) {
+      if (CometConf.COMET_STRICT_FALLBACK_REASONS.get(op.conf)) {
+        throw new IllegalStateException(
+          s"Comet did not convert ${op.nodeName} but recorded no fallback reason on the " +
+            "operator or any of its expressions. Add a withFallbackReason call stating why " +
+            s"conversion failed. Operator:\n$op")
+      }
+      withFallbackReason(op, s"${op.nodeName} is not supported")
+    }
+  }
+
   /**
    * 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

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

@@ -21,7 +21,6 @@ package org.apache.comet.serde
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, BloomFilterMightContain}
 
-import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
 import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
 
 object CometBloomFilterMightContain extends CometExpressionSerde[BloomFilterMightContain] {
@@ -45,7 +44,6 @@ object CometBloomFilterMightContain extends CometExpressionSerde[BloomFilterMigh
           .setBloomFilterMightContain(builder)
           .build())
     } else {
-      withFallbackReason(expr, bloomFilter, value)
       None
     }
   }

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

@@ -22,13 +22,13 @@ package org.apache.comet.serde
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
 
 import org.apache.comet.serde.ExprOuterClass.Expr
-import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFallbackReason, scalarFunctionExprToProto}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProto}
 
 /** Serde for scalar function. */
 case class CometScalarFunction[T <: Expression](name: String) extends CometExpressionSerde[T] {
   override def convert(expr: T, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = {
     val childExpr = expr.children.map(exprToProtoInternal(_, inputs, binding))
     val optExpr = scalarFunctionExprToProto(name, childExpr: _*)
-    optExprWithFallbackReason(optExpr, expr, expr.children: _*)
+    optExpr
   }
 }

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

@@ -22,7 +22,6 @@ package org.apache.comet.serde
 import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Descending, NullsFirst, NullsLast, SortOrder}
 
 import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
 import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
 
 object CometSortOrder extends CometExpressionSerde[SortOrder] {
@@ -66,7 +65,6 @@ object CometSortOrder extends CometExpressionSerde[SortOrder] {
           .setSortOrder(sortOrderBuilder)
           .build())
     } else {
-      withFallbackReason(expr, expr.child)
       None
     }
   }