fix test failures

Author: mbutrovich

spark/src/main/scala/org/apache/comet/planner/CometPlanner.scala

@@ -89,7 +89,7 @@ case class CometPlanner(session: SparkSession) extends Rule[SparkPlan] with Logg
     // Phase 1 must run before BroadcastConsumerIndex.build: Phase 1's generic-exec prediction
     // reads children's LIKELY_COMET tags (post-order walk), and the index reads BHJ tags to
     // decide which broadcasts have a Comet consumer.
-    val annotated1 = phase1LikelyComet(
+    val annotated1: SparkPlan = phase1LikelyComet(
       prepared,
       PlanningContext(
         session = session,
@@ -107,8 +107,9 @@ case class CometPlanner(session: SparkSession) extends Rule[SparkPlan] with Logg
 
     val annotated2 = phase2Decision(annotated1, context)
     val emitted = phase3Emit(annotated2, context)
-    val reverted = revertOrphanedBroadcasts(emitted)
-    val cleaned = cleanupLogicalLinks(reverted)
+    val broadcastsReverted = revertBroadcastsWithoutCometConsumer(emitted)
+    val shufflesReverted = revertRedundantColumnarShuffle(broadcastsReverted)
+    val cleaned = cleanupLogicalLinks(shufflesReverted)
     val blocked = convertBlocks(cleaned)
     val finalPlan = postPass(blocked, context)
 
@@ -179,7 +180,7 @@ case class CometPlanner(session: SparkSession) extends Rule[SparkPlan] with Logg
    * Shuffle doesn't need the equivalent revert because a Spark parent with a Comet columnar
    * shuffle child is handled naturally by Spark's transition insertion.
    */
-  private def revertOrphanedBroadcasts(plan: SparkPlan): SparkPlan = {
+  private def revertBroadcastsWithoutCometConsumer(plan: SparkPlan): SparkPlan = {
     if (CometConf.COMET_EXEC_BROADCAST_FORCE_ENABLED.get()) {
       return plan
     }
@@ -199,6 +200,50 @@ case class CometPlanner(session: SparkSession) extends Rule[SparkPlan] with Logg
     out
   }
 
+  /**
+   * Revert a `CometShuffleExchangeExec` with `CometColumnarShuffle` whose parent and child are
+   * both non-Comet `HashAggregateExec` / `ObjectHashAggregateExec` back to the original Spark
+   * `ShuffleExchangeExec`. Mirrors the legacy `revertRedundantColumnarShuffle` (PR #4010): the
+   * partial-final-aggregate pattern where both aggregates fall back to Spark would otherwise keep
+   * a columnar shuffle between them, adding row->arrow->shuffle->arrow->row conversion with no
+   * Comet consumer on either side.
+   *
+   * Phase 1's optimistic-true prediction for shuffles allows the legitimate
+   * `Sort-over-Spark-leaf` pattern to convert (the shuffle does row->arrow at exchange time). The
+   * same optimism produces the redundant pattern when both ends remain Spark, which this pass
+   * cleans up. Narrow match on aggregate-shuffle-aggregate keeps the intervention surgical; other
+   * Spark-Comet-Spark sandwiches are handled by `revertBroadcastsWithoutCometConsumer` or Spark's
+   * transition insertion.
+   */
+  private def revertRedundantColumnarShuffle(plan: SparkPlan): SparkPlan = {
+    def isAggregate(p: SparkPlan): Boolean =
+      p.isInstanceOf[org.apache.spark.sql.execution.aggregate.HashAggregateExec] ||
+        p.isInstanceOf[org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec]
+
+    def isRedundantShuffle(child: SparkPlan): Boolean = child match {
+      case s: CometShuffleExchangeExec =>
+        s.shuffleType == org.apache.spark.sql.comet.execution.shuffle.CometColumnarShuffle &&
+        isAggregate(s.child)
+      case _ => false
+    }
+
+    var reverted = 0
+    val out = plan.transform {
+      case op if isAggregate(op) && op.children.exists(isRedundantShuffle) =>
+        val newChildren = op.children.map {
+          case s: CometShuffleExchangeExec
+              if s.shuffleType == org.apache.spark.sql.comet.execution.shuffle.CometColumnarShuffle
+                && isAggregate(s.child) =>
+            reverted += 1
+            s.originalPlan.withNewChildren(Seq(s.child))
+          case other => other
+        }
+        op.withNewChildren(newChildren)
+    }
+    if (reverted > 0) logDebug(s"CometPlanner: reverted $reverted redundant columnar shuffles")
+    out
+  }
+
   private def isNativeCompatible(node: SparkPlan): Boolean =
     node.isInstanceOf[CometNativeExec] || node.getTagValue(CometTags.NATIVE_OP).isDefined
 

spark/src/main/scala/org/apache/comet/planner/phases/Phase1LikelyComet.scala

@@ -99,23 +99,18 @@ object Phase1LikelyComet extends Logging {
     case _: BatchScanExec =>
       CometConf.COMET_NATIVE_SCAN_ENABLED.get(conf)
 
-    case s: ShuffleExchangeExec =>
-      // A shuffle's LIKELY_COMET reflects whether its data source can provide native input.
-      // Optimistic-true broke upward propagation: a Spark partial HashAgg below a shuffle could
-      // not flag the shuffle as non-likely, which made the final HashAgg above read a stale-true
-      // child, which made Phase 2 see `parentLikely=true` on the shuffle and convert it. The
-      // legacy rule lived with this and added `revertRedundantColumnarShuffle` (PR #4010) as a
-      // post-pass; the planner avoids the conversion in the first place by predicting
-      // accurately. Phase 3 still re-checks at emit time. S2C-eligible leaves count because
-      // Phase 2 wraps them in `CometSparkToColumnarExec`.
-      s.children.exists(c => childCanProvideNativeInput(c, conf))
+    case _: ShuffleExchangeExec =>
+      // Optimistic: shuffles can absorb non-native children via `CometColumnarShuffle` (the
+      // shuffle's row->arrow conversion runs at exchange time). Tying the prediction to the
+      // child here would prevent the legacy `Sort over Spark LocalTableScan via Comet shuffle`
+      // shape from working. The redundant `HashAgg(JVM) -> CometShuffle -> HashAgg(JVM)` case
+      // that this optimism would otherwise create is handled by `revertRedundantColumnarShuffle`
+      // as a post-pass (mirrors PR #4010). Phase 3 re-checks at emit time.
+      true
 
-    case b: BroadcastExchangeExec =>
-      // Same shape as shuffle: a broadcast over Spark-only data cannot itself go native, so
-      // Phase 1 must report that honestly for parent BHJs' children-OK checks to be correct.
-      // Phase 2's BroadcastConsumerIndex still gates conversion on a downstream Comet BHJ
-      // wanting this broadcast; this only changes whether the broadcast is a candidate at all.
-      b.children.exists(c => childCanProvideNativeInput(c, conf))
+    case _: BroadcastExchangeExec =>
+      // Same as shuffle. Phase 3 re-checks at emit time.
+      true
 
     // AQE stage re-entry: a prior CometPlanner pass converted an exchange, AQE materialized it
     // and wrapped it in a query stage. Phase 3 re-emits the stage itself as a Comet-compatible

spark/src/test/scala/org/apache/comet/planner/CometPlannerSuite.scala

@@ -274,7 +274,7 @@ class CometPlannerSuite extends CometTestBase {
 
   // --- Demand-aware shuffle contract (translation of revert-pass tests) ----------------------
 
-  test("CometPlanner should not emit Comet shuffle between Spark aggregates") {
+  test("CometPlanner should revert redundant Comet shuffle between Spark aggregates") {
     withTempView("test_data") {
       createTestDataFrame.createOrReplaceTempView("test_data")
 
@@ -284,11 +284,13 @@ class CometPlannerSuite extends CometTestBase {
       assert(countOperators(sparkPlan, classOf[ShuffleExchangeExec]) == 1)
       assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2)
 
-      // Disable partial aggregate so both aggregates fall back to Spark JVM. The planner's
-      // Phase 2 demand-aware rule should emit Passthrough for the shuffle (selfLikely=true but
-      // parent/child both have LIKELY_COMET=false), so the shuffle stays Spark with no revert
-      // pass needed. Minimal reproducer for the redundant-shuffle pattern that PR #4010 fixed
-      // for the legacy rule via revertRedundantColumnarShuffle.
+      // Disable partial aggregate so both aggregates fall back to Spark JVM. Phase 1's
+      // optimistic shuffle prediction lets Phase 2 speculatively convert the shuffle (needed
+      // for the legitimate Sort-over-Spark-leaf shape). The post-Phase-3
+      // `revertRedundantColumnarShuffle` pass then detects the resulting
+      // `HashAgg(JVM) -> CometColumnarShuffle -> HashAgg(JVM)` pattern and reverts the shuffle
+      // to plain Spark, avoiding the row->arrow->shuffle->arrow->row round-trip with no Comet
+      // consumer on either side. Mirrors PR #4010.
       withSQLConf(
         CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key -> "false",
         CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {