Add suites to yml

Author: karuppayya

.github/workflows/pr_build_linux.yml

@@ -321,6 +321,7 @@ jobs:
               org.apache.spark.CometPluginsUnifiedModeOverrideSuite
               org.apache.comet.rules.CometScanRuleSuite
               org.apache.comet.rules.CometExecRuleSuite
+              org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite
               org.apache.spark.sql.CometTPCDSQuerySuite
               org.apache.spark.sql.CometTPCDSQueryTestSuite
               org.apache.spark.sql.CometTPCHQuerySuite

.github/workflows/pr_build_macos.yml

@@ -163,6 +163,7 @@ jobs:
               org.apache.spark.CometPluginsUnifiedModeOverrideSuite
               org.apache.comet.rules.CometScanRuleSuite
               org.apache.comet.rules.CometExecRuleSuite
+              org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite
               org.apache.spark.sql.CometTPCDSQuerySuite
               org.apache.spark.sql.CometTPCDSQueryTestSuite
               org.apache.spark.sql.CometTPCHQuerySuite

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

@@ -465,8 +465,10 @@ object CometConf extends ShimCometConf {
           "shuffle still requires at least one R2C at the shuffle boundary. " +
           "Only effective when spark.comet.exec.transitionRevert.enabled is true.")
       .intConf
-      .checkValue(_ >= 2, "Must be >= 2. A reverted stage still requires at least one " +
-        "R2C at the columnar shuffle boundary.")
+      .checkValue(
+        _ >= 2,
+        "Must be >= 2. A reverted stage still requires at least one " +
+          "R2C at the columnar shuffle boundary.")
       .createWithDefault(2)
 
   val COMET_EXEC_SHUFFLE_COMPRESSION_CODEC: ConfigEntry[String] =

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

@@ -107,9 +107,8 @@ class CometSparkSessionExtensions
     override def preColumnarTransitions: Rule[SparkPlan] = CometExecRule(session)
 
     override def postColumnarTransitions: Rule[SparkPlan] = {
-      val rules = Seq(
-        EliminateRedundantTransitions(session),
-        RevertNativeForTransitionHeavyStages(session))
+      val rules =
+        Seq(EliminateRedundantTransitions(session), RevertNativeForTransitionHeavyStages(session))
       plan => rules.foldLeft(plan) { case (p, rule) => rule(p) }
     }
   }

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

@@ -34,7 +34,6 @@ import org.apache.comet.CometConf
  * Reverts a query stage to Spark row-based execution when it has too many columnar-to-row (C2R)
  * transitions. Each C2R indicates Comet could not keep execution columnar and had to fall back.
  * With columnar shuffle enabled, each C2R implies a corresponding R2C round-trip.
- *
  */
 case class RevertNativeForTransitionHeavyStages(session: SparkSession)
     extends Rule[SparkPlan]
@@ -66,15 +65,16 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession)
   }
 
   private def applyForNonAQE(plan: SparkPlan): SparkPlan = {
-    plan.transformUp {
-      case exchange: ShuffleExchangeLike =>
-        revertStageIfNeeded(exchange.child, exchange.supportsColumnar)
-          .map(reverted => exchange.withNewChildren(Seq(reverted)))
-          .getOrElse(exchange)
+    plan.transformUp { case exchange: ShuffleExchangeLike =>
+      revertStageIfNeeded(exchange.child, exchange.supportsColumnar)
+        .map(reverted => exchange.withNewChildren(Seq(reverted)))
+        .getOrElse(exchange)
     }
   }
 
-  /** Reverts the stage if C2R count exceeds threshold. Wraps in R2C if exchange needs columnar. */
+  /**
+   * Reverts the stage if C2R count exceeds threshold. Wraps in R2C if exchange needs columnar.
+   */
   private def revertStageIfNeeded(
       stagePlan: SparkPlan,
       outputColumnar: Boolean): Option[SparkPlan] = {
@@ -94,7 +94,6 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession)
     Some(result)
   }
 
-
   /** Counts C2R transitions within this stage, stopping at stage boundaries. */
   private[rules] def countTransitions(plan: SparkPlan): Int = {
     var count = 0

spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala

@@ -19,13 +19,13 @@
 
 package org.apache.comet.rules
 
+import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
 
 import org.apache.comet.CometConf
-import org.apache.spark.sql.CometTestBase
 
 class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
 
@@ -77,7 +77,8 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
       CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "10") {
       withTempView("test_data") {
         spark.range(10).toDF("id").createOrReplaceTempView("test_data")
-        val sparkPlan = createSparkPlan("SELECT id, id * 2 as doubled FROM test_data WHERE id > 5")
+        val sparkPlan =
+          createSparkPlan("SELECT id, id * 2 as doubled FROM test_data WHERE id > 5")
         val cometPlan = applyCometExecRule(sparkPlan)
 
         val rule = RevertNativeForTransitionHeavyStages(spark)
@@ -141,7 +142,8 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
         val rule = RevertNativeForTransitionHeavyStages(spark)
         val reverted = rule.revertToSpark(cometPlan)
 
-        assert(countCometExecs(reverted) == 0,
+        assert(
+          countCometExecs(reverted) == 0,
           s"Should have no CometExec nodes after revert, plan:\n${reverted.treeString}")
       }
     }
@@ -162,7 +164,8 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
         val reverted = rule.revertToSpark(cometPlan)
 
         // Reverted plan should have same output schema
-        assert(reverted.output.map(_.name) == cometPlan.output.map(_.name),
+        assert(
+          reverted.output.map(_.name) == cometPlan.output.map(_.name),
           "Output schema should be preserved after revert")
       }
     }
@@ -183,7 +186,8 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
 
         val rule = RevertNativeForTransitionHeavyStages(spark)
         val result = rule.revertToSpark(cometPlan)
-        assert(countCometExecs(result) == 0,
+        assert(
+          countCometExecs(result) == 0,
           s"All CometExec should be reverted. Plan:\n${result.treeString}")
       }
     }
@@ -205,14 +209,14 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
         if (cometShuffles.nonEmpty) {
           val rule = RevertNativeForTransitionHeavyStages(spark)
           val reverted = rule.revertToSpark(cometPlan)
-          val remainingCometShuffles = reverted.collect {
-            case s: CometShuffleExchangeExec => s
+          val remainingCometShuffles = reverted.collect { case s: CometShuffleExchangeExec =>
+            s
           }
-          assert(remainingCometShuffles.isEmpty,
+          assert(
+            remainingCometShuffles.isEmpty,
             "CometShuffleExchangeExec should be reverted to ShuffleExchangeExec")
           val sparkShuffles = reverted.collect { case s: ShuffleExchangeExec => s }
-          assert(sparkShuffles.nonEmpty,
-            "Should have ShuffleExchangeExec after revert")
+          assert(sparkShuffles.nonEmpty, "Should have ShuffleExchangeExec after revert")
         }
       }
     }
@@ -228,17 +232,17 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
       "spark.sql.adaptive.enabled" -> "false") {
 
       withTempView("test_data") {
-        spark.range(10).selectExpr("id", "id % 3 as grp")
+        spark
+          .range(10)
+          .selectExpr("id", "id % 3 as grp")
           .createOrReplaceTempView("test_data")
-        val sparkPlan = createSparkPlan(
-          "SELECT grp, count(*) FROM test_data GROUP BY grp")
+        val sparkPlan = createSparkPlan("SELECT grp, count(*) FROM test_data GROUP BY grp")
         val cometPlan = applyCometExecRule(sparkPlan)
 
         // With high threshold, the non-AQE path should not revert anything
         val rule = RevertNativeForTransitionHeavyStages(spark)
         val result = rule.apply(cometPlan)
-        assert(result eq cometPlan,
-          "Non-AQE path should not revert when below threshold")
+        assert(result eq cometPlan, "Non-AQE path should not revert when below threshold")
       }
     }
   }
@@ -259,10 +263,12 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
         val pairs = rule.countTransitions(cometPlan)
         val result = rule.apply(cometPlan)
         if (pairs <= 2) {
-          assert(result eq cometPlan,
+          assert(
+            result eq cometPlan,
             s"Plan with $pairs pairs should NOT be reverted at threshold 2")
         } else {
-          assert(countCometExecs(result) == 0,
+          assert(
+            countCometExecs(result) == 0,
             s"Plan with $pairs pairs should be reverted at threshold 2")
         }
       }