simplify code

apache · Jul 1, 2024 · 07ca195 · 07ca195
1 parent c6eaae5
commit 07ca195
Showing 1 changed file with 5 additions and 10 deletions.
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -202,7 +202,6 @@ class CometSparkSessionExtensions
     private val execRule = CometExecRule(session)
 
     def apply(plan: SparkPlan): SparkPlan = {
-      var needToReplan = false
       val cometPlan = execRule.apply(plan)
       if (CometConf.COMET_CBO_ENABLED.get()) {
         // simple heuristic to avoid moving from Spark execution to Comet execution just
@@ -211,17 +210,13 @@ class CometSparkSessionExtensions
         cometPlan match {
           case CometSortExec(_, _, _, e: CometShuffleExchangeExec, _)
               if !e.child.supportsColumnar =>
-            needToReplan = true
-            // fall back for sort
-            plan.setTagValue(CometExplainInfo.CBO_FALLBACK, "avoid move to Comet just for sort")
-            // fall back for exchange as well
-            plan.children.head
-              .setTagValue(CometExplainInfo.CBO_FALLBACK, "avoid move to Comet just for sort")
+            // fall back for sort amd exchange operators
+            val fallbackReason = "avoid move to Comet just for sort"
+            plan.setTagValue(CometExplainInfo.CBO_FALLBACK, fallbackReason)
+            plan.children.head.setTagValue(CometExplainInfo.CBO_FALLBACK, fallbackReason)
+            return execRule.apply(plan)
           case _ =>
         }
-        if (needToReplan) {
-          return execRule.apply(plan)
-        }
       }
       cometPlan
     }