From 07ca19558d1c48738bf1560ab1e893f48992c8bb Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 1 Jul 2024 11:29:26 -0600
Subject: [PATCH] simplify code

---
 .../comet/CometSparkSessionExtensions.scala       | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index b279bc014..09583dc7a 100644
--- 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
     }