Optimize implementation of getAggregateRawMetrics in core-tools

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/AppSparkMetricsAnalyzer.scala

@@ -25,7 +25,7 @@ import com.nvidia.spark.rapids.tool.profiling._
 
 import org.apache.spark.sql.rapids.tool.{AppBase, ToolUtils}
 import org.apache.spark.sql.rapids.tool.profiling.ApplicationInfo
-import org.apache.spark.sql.rapids.tool.store.{AccumInfo, AccumMetaRef, AccumNameRef}
+import org.apache.spark.sql.rapids.tool.store.{AccumInfo, AccumMetaRef}
 
 /**
  * Does analysis on the DataFrames from object of AppBase.
@@ -85,10 +85,10 @@ class AppSparkMetricsAnalyzer(app: AppBase) extends AppAnalysisBase(app) {
       } else {
         val jobAggAccumulator = new AggAccumHelper()
         val perJobRec = jobAggAccumulator.accumPerJob(
-          jc.stageIds.filter(stageLevelSparkMetrics(index).contains)
-            .map { stageId =>
+          jc.stageIds.collect {
+            case stageId if stageLevelSparkMetrics(index).contains(stageId) =>
               stageLevelSparkMetrics(index)(stageId)
-            })
+          })
         if (perJobRec.isEmptyAggregates) {
           None
         } else {
@@ -178,10 +178,10 @@ class AppSparkMetricsAnalyzer(app: AppBase) extends AppAnalysisBase(app) {
         // TODO: Should we only consider successful tasks?
         val sqlAggAccumulator = new AggAccumHelper()
         val preSqlRec = sqlAggAccumulator.accumPerSQL(
-          stagesInSQL.filter(stageLevelSparkMetrics(index).contains)
-            .map { stageId =>
+          stagesInSQL.collect {
+            case stageId if stageLevelSparkMetrics(index).contains(stageId) =>
               stageLevelSparkMetrics(index)(stageId)
-            })
+          })
         if (preSqlRec.isEmptyAggregates) {
           None
         } else {
@@ -322,20 +322,21 @@ class AppSparkMetricsAnalyzer(app: AppBase) extends AppAnalysisBase(app) {
         app.asInstanceOf[ApplicationInfo].planMetricProcessor
     }
     val zeroAccumProfileResults =
-      AccumProfileResults(0, 0, AccumMetaRef(0L, AccumNameRef("")), 0L, 0L, 0L, 0L)
-
+      AccumProfileResults(0, 0, AccumMetaRef.EMPTY_ACCUM_META_REF, 0L, 0L, 0L, 0L)
+    val emptyNodeNames = Seq.empty[String]
+    val emptyDiagnosticMetrics = HashMap.empty[String, AccumProfileResults]
     // TODO: this has stage attempts. we should handle different attempts
     app.stageManager.getAllStages.map { sm =>
       // TODO: Should we only consider successful tasks?
       val tasksInStage = app.taskManager.getTasks(sm.stageInfo.stageId,
         sm.stageInfo.attemptNumber())
       // count duplicate task attempts
       val numTasks = tasksInStage.size
-      val nodeNames = sqlAnalyzer.stageToNodeNames.
-        getOrElse(sm.stageInfo.stageId, Seq.empty[String])
-      val diagnosticMetricsMap = sqlAnalyzer.stageToDiagnosticMetrics.
-        getOrElse(sm.stageInfo.stageId, HashMap.empty[String, AccumProfileResults]).
-        withDefaultValue(zeroAccumProfileResults)
+      val nodeNames = sqlAnalyzer.stageToNodeNames.getOrElse(sm.stageInfo.stageId, emptyNodeNames)
+      val diagnosticMetricsMap =
+        sqlAnalyzer.stageToDiagnosticMetrics
+          .getOrElse(sm.stageInfo.stageId, emptyDiagnosticMetrics)
+          .withDefaultValue(zeroAccumProfileResults)
       val srTotalBytesMetrics =
         AppSparkMetricsAnalyzer.getStatistics(tasksInStage.map(_.sr_totalBytesRead))
 
@@ -450,8 +451,6 @@ class AppSparkMetricsAnalyzer(app: AppBase) extends AppAnalysisBase(app) {
         perStageRec.srTotalBytesReadSum,
         perStageRec.swBytesWrittenSum,
         perStageRec.swRecordsWrittenSum,
-        // Leave this timeUnit in NanoSeconds so that it will be more accurate when we take
-        // aggregates on higher levels (i.e., SQL/Job)
         perStageRec.swWriteTimeSum)
       stageLevelSparkMetrics(index).put(sm.stageInfo.stageId, stageRow)
     }
@@ -477,12 +476,4 @@ object AppSparkMetricsAnalyzer  {
       StatisticsMetrics(sortedArr.head, med, sortedArr(len - 1), sortedArr.sum)
     }
   }
-
-  def maxWithEmptyHandling(arr: Iterable[Long]): Long = {
-    if (arr.isEmpty) {
-      0L
-    } else {
-      arr.max
-    }
-  }
 }

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/AggAccumHelper.scala

@@ -26,11 +26,6 @@ import org.apache.spark.sql.rapids.tool.store.TaskModel
  * a parellel processor can be used to split the iterables without changing the caller side.
  */
 class AggAccumHelper {
-  private def initializeRecord(rec: TaskMetricsAccumRec, iterable: Iterable[Any]): Unit = {
-    if (iterable.isEmpty) { // Reset aggregate fields for empty collections
-      rec.resetFields()
-    }
-  }
 
   private def accumCachedRecords[R <: TaskMetricsAccumRec](
       stageRecords: Iterable[StageAggTaskMetricsProfileResult],
@@ -45,22 +40,19 @@ class AggAccumHelper {
 
   def accumPerStage(taskRecords: Iterable[TaskModel]): TaskMetricsAccumRec = {
     val resRec = createStageAccumRecord()
-    initializeRecord(resRec, taskRecords)
     taskRecords.foreach(resRec.addRecord)
     resRec.finalizeAggregation()
     resRec
   }
 
   def accumPerSQL(stageRecords: Iterable[StageAggTaskMetricsProfileResult]): SQLAggAccum = {
     val resRec = SQLAggAccum()
-    initializeRecord(resRec, stageRecords)
     accumCachedRecords(stageRecords, resRec)
     resRec
   }
 
   def accumPerJob(stageRecords: Iterable[StageAggTaskMetricsProfileResult]): JobAggAccum = {
     val resRec = JobAggAccum()
-    initializeRecord(resRec, stageRecords)
     accumCachedRecords(stageRecords, resRec)
     resRec
   }

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/StageAggPhoton.scala

@@ -39,11 +39,15 @@ case class StageAggPhoton(
     // the peakMemValues.
     swWriteTimeSum = 0
     peakExecutionMemoryMax = 0
-    if (shuffleWriteValues.nonEmpty) {
-      swWriteTimeSum = TimeUnit.NANOSECONDS.toMillis(shuffleWriteValues.sum)
-    }
-    if (peakMemValues.nonEmpty) {
-      peakExecutionMemoryMax = peakMemValues.max
+    if (!isEmptyAggregates) {
+      // Re-calculate the photon specific fields only if the accumulator has tasks.
+      // Otherwise, leave it as 0.
+      if (shuffleWriteValues.nonEmpty) {
+        swWriteTimeSum = TimeUnit.NANOSECONDS.toMillis(shuffleWriteValues.sum)
+      }
+      if (peakMemValues.nonEmpty) {
+        peakExecutionMemoryMax = peakMemValues.max
+      }
     }
     super.finalizeAggregation()
   }

core/src/main/scala/com/nvidia/spark/rapids/tool/analysis/util/TaskMetricsAccumRec.scala

@@ -63,12 +63,17 @@ class TaskMetricsAccumRec {
    */
   def isEmptyAggregates: Boolean = numTasks == 0
 
+  /**
+   * Reset all fields to 0. This is used to reset the fields when the Task iterator is empty.
+   * When the iterator is empty, then fields such as "max" should be reset to 0.
+   */
   def resetFields(): Unit = {
     durationMax = 0
     durationMin = 0
     peakExecutionMemoryMax = 0
     resultSizeMax = 0
   }
+
   def addRecord(rec: TaskModel): Unit = {
     numTasks += 1
     // SumFields
@@ -102,6 +107,7 @@ class TaskMetricsAccumRec {
     // Min Fields
     durationMin = math.min(durationMin, rec.duration)
   }
+
   def addRecord(rec: StageAggTaskMetricsProfileResult): Unit = {
     // Sums
     numTasks += rec.numTasks
@@ -143,5 +149,9 @@ class TaskMetricsAccumRec {
    */
   def finalizeAggregation(): Unit = {
     durationAvg = ToolUtils.calculateAverage(durationSum, numTasks, 1)
+    if (numTasks < 1) {
+      // number of tasks is 0, then we should reset fields such as (max, min) to 0.
+      resetFields()
+    }
   }
 }

core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumMetaRef.scala

@@ -27,6 +27,7 @@ case class AccumMetaRef(id: Long, name: AccumNameRef) {
 }
 
 object AccumMetaRef {
+  val EMPTY_ACCUM_META_REF: AccumMetaRef = new AccumMetaRef(0L, AccumNameRef.EMPTY_ACC_NAME_REF)
   def apply(id: Long, name: Option[String]): AccumMetaRef =
     new AccumMetaRef(id, AccumNameRef.getOrCreateAccumNameRef(name))
 }

core/src/main/scala/org/apache/spark/sql/rapids/tool/store/AccumNameRef.scala

@@ -42,7 +42,7 @@ case class AccumNameRef(value: String) {
 object AccumNameRef {
   // Dummy AccNameRef to represent None accumulator names. This is an optimization to avoid
   // storing an option[string] for all accumulable names which leads to "get-or-else" everywhere.
-  private val EMPTY_ACC_NAME_REF: AccumNameRef = new AccumNameRef("N/A")
+  val EMPTY_ACC_NAME_REF: AccumNameRef = new AccumNameRef("N/A")
   // A global table to store reference to all accumulator names. The map is accessible by all
   // threads (different applications) running in parallel. This avoids duplicate work across
   // different threads.