Address review comments

apache · Apr 11, 2024 · bdf6b6d · bdf6b6d
1 parent dae33d9
commit bdf6b6d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 15 deletions.
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -484,7 +484,6 @@ class CometSparkSessionExtensions
               CometTakeOrderedAndProjectExec
                 .isSupported(s)
                 ._1 =>
-          // TODO: support offset for Spark 3.4
           val (newOp, info) = QueryPlanSerde.operator2Proto(s)
           newOp match {
             case Some(nativeOp) =>
@@ -691,9 +690,7 @@ class CometSparkSessionExtensions
         case op =>
           // An operator that is not supported by Comet
           op match {
-            case b: CometExec => b
-            case b: CometBroadcastExchangeExec => b
-            case b: CometShuffleExchangeExec => b
+            case _: CometExec | _: CometBroadcastExchangeExec | _: CometShuffleExchangeExec => op
             case o =>
               opWithInfo(o, CometExplainInfo(s"${o.nodeName} is not supported"))
           }

diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala
@@ -47,13 +47,13 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
   private def extensionInfo(plan: SparkPlan): mutable.Seq[String] = {
     var info = mutable.Seq[String]()
     val sorted = sortup(plan)
-    sorted.foreach(p => {
+    sorted.foreach { p =>
       val s =
-        getActualPlan(p).getTagValue(CometExplainInfo.EXTENSION_INFO).map(t => t).getOrElse("")
+        getActualPlan(p).getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse("")
       if (s.nonEmpty) {
         info = info :+ s
       }
-    })
+    }
     info
   }
 
@@ -65,18 +65,14 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator {
       val s = traversed.dequeue()
       ordered += s
       if (s.innerChildren.nonEmpty) {
-        s.innerChildren.foreach(c => {
-          c match {
-            case _: SparkPlan => traversed.enqueue(getActualPlan(c.asInstanceOf[SparkPlan]))
-            case _ =>
-          }
-          ()
-        })
+        s.innerChildren.foreach {
+          case c@(_: SparkPlan) => traversed.enqueue(getActualPlan(c))
+          case _ =>
+        }
       }
       if (s.children.nonEmpty) {
         s.children.foreach(c => {
           traversed.enqueue(getActualPlan(c))
-          ()
         })
       }
     }