optimize checks

Signed-off-by: Niranjan Artal <[email protected]>
NVIDIA · Nov 3, 2023 · 3573d3e · 3573d3e
1 parent a94e898
commit 3573d3e
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/.../src/main/scala/org/apache/spark/sql/rapids/tool/qualification/QualificationAppInfo.scala b/.../src/main/scala/org/apache/spark/sql/rapids/tool/qualification/QualificationAppInfo.scala
@@ -201,12 +201,12 @@ class QualificationAppInfo(
   private def checkStageIdInExec(prev: Option[ExecInfo],
       execInfo: ExecInfo, next: Option[ExecInfo]): (Seq[(Int, ExecInfo)], Option[ExecInfo]) = {
     val associatedStages = {
-      if (execInfo.stages.size >= 1) {
+      if (execInfo.stages.nonEmpty) {
         execInfo.stages.toSeq
       } else {
-        if (prev.exists(_.stages.size >= 1)) {
+        if (prev.exists(_.stages.nonEmpty)) {
           prev.flatMap(_.stages.headOption).toSeq
-        } else if (next.nonEmpty) {
+        } else if (next.exists(_.stages.nonEmpty)) {
           next.flatMap(_.stages.headOption).toSeq
         } else {
           // we don't know what stage its in or its duration
@@ -223,7 +223,6 @@ class QualificationAppInfo(
     }
   }
 
-
   private def getStageToExec(execInfos: Seq[ExecInfo]): (Map[Int, Seq[ExecInfo]], Seq[ExecInfo]) = {
     val execsWithoutStages = new ArrayBuffer[ExecInfo]()