Qualification tool: Enhance mapping of Execs to stages

core/src/main/scala/org/apache/spark/sql/rapids/tool/qualification/QualificationAppInfo.scala

@@ -200,22 +200,79 @@ class QualificationAppInfo(
 
   private def getStageToExec(execInfos: Seq[ExecInfo]): (Map[Int, Seq[ExecInfo]], Seq[ExecInfo]) = {
     val execsWithoutStages = new ArrayBuffer[ExecInfo]()
-    val perStageSum = execInfos.flatMap { execInfo =>
-      if (execInfo.stages.size > 1) {
-        execInfo.stages.map((_, execInfo))
-      } else if (execInfo.stages.size < 1) {
-        // we don't know what stage its in or its duration
-        logDebug(s"No stage associated with ${execInfo.exec} " +
-          s"so speedup factor isn't applied anywhere.")
-        execsWithoutStages += execInfo
-        Seq.empty
-      } else {
-        Seq((execInfo.stages.head, execInfo))
-      }
-    }.groupBy(_._1).map { case (stage, execInfos) =>
+
+    // This is to get the mapping between stageId and execs. This is primarily done based on
+    // accumulatorId. If an Exec has some metrics generated, then an accumulatorId will be
+    // generated for that Exec. This accumulatorId is used to get the stageId. If an Exec
+    // doesn't have any metrics, then we will try to get the stageId by looking at the
+    // neighbor Execs. If either of the neighbor Execs has a stageId, then we will use that
+    // to assign the same stageId to the current Exec as it's most likely that the current
+    // Exec is part of the same stage as the neighbor Exec.
+    val execInfosInOrder = execInfos.reverse
+    val execInfosCombined = Iterator(Seq(execInfosInOrder.head)) ++
+      execInfosInOrder.sliding(3) ++ Iterator(Seq(execInfosInOrder.last))
+    val perStageSum = execInfosCombined.map {
+      case Seq(prev, execInfo, next) =>
+        val associatedStages =
+          if (execInfo.stages.size > 1) execInfo.stages.toSeq
+          else if (execInfo.stages.size < 1) {
+            val prevStages = if (prev.stages.size > 1) prev.stages else Seq.empty
+            val nextStages = if (next.stages.size > 1) next.stages else Seq.empty
+            val singlePrevStage = prev.stages.headOption.toSeq
+            val singleNextStage = next.stages.headOption.toSeq
+            val dedupedStages = Set(singlePrevStage ++ singleNextStage ++ prevStages ++ nextStages).
+              flatten.toSeq
+            if (dedupedStages.nonEmpty) {
+              dedupedStages
+            } else {
+              // we don't know what stage its in or its duration
+              logDebug(s"No stage associated with ${execInfo.exec} " +
+                s"so speedup factor isn't applied anywhere.")
+              execsWithoutStages += execInfo
+              Seq.empty
+            }
+          } else {
+            Seq(execInfo.stages.head)
+          }
+        associatedStages.map((_, execInfo))
+
+      case Seq(prev, execInfo) =>
+        val associatedStages =
+          if (execInfo.stages.size > 1) execInfo.stages.toSeq
+          else if (execInfo.stages.size < 1) {
+            val prevStages = if (prev.stages.size > 1) prev.stages.toSeq else Seq.empty
+            val singlePrevStage = prev.stages.headOption.toSeq
+            val dedupedStages = Set(singlePrevStage ++ prevStages).flatten.toSeq
+            if (dedupedStages.nonEmpty) {
+              dedupedStages
+            } else {
+              // we don't know what stage its in or its duration
+              logDebug(s"No stage associated with ${execInfo.exec} " +
+                s"so speedup factor isn't applied anywhere.")
+              execsWithoutStages += execInfo
+              Seq.empty
+            }
+          } else {
+            Seq(execInfo.stages.head)
+          }
+        associatedStages.map((_, execInfo))
+
+      case Seq(execInfo) =>
+        val associatedStages =
+          if (execInfo.stages.size > 1) execInfo.stages.toSeq
+          else if (execInfo.stages.size < 1) {
+            execsWithoutStages += execInfo
+            Seq.empty
+          } else {
+            Seq(execInfo.stages.head)
+          }
+        associatedStages.map((_, execInfo))
+
+      case _ => Seq.empty
+    }.toSeq.flatten.groupBy(_._1).map { case (stage, execInfos) =>
       (stage, execInfos.map(_._2))
     }
-    (perStageSum, execsWithoutStages.toSeq)
+    (perStageSum, execsWithoutStages)
   }
 
   private def flattenedExecs(execs: Seq[ExecInfo]): Seq[ExecInfo] = {

core/src/test/resources/QualificationExpectations/jdbc_expectation.csv

@@ -1,2 +1,2 @@
 App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Estimated GPU Time Saved,SQL DF Duration,SQL Dataframe Task Duration,App Duration,GPU Opportunity,Executor CPU Time Percent,SQL Ids with Failures,Unsupported Read File Formats and Types,Unsupported Write Data Format,Complex Types,Nested Complex Types,Potential Problems,Longest SQL Duration,NONSQL Task Duration Plus Overhead,Unsupported Task Duration,Supported SQL DF Task Duration,Task Speedup Factor,App Duration Estimated,Unsupported Execs,Unsupported Expressions,Estimated Job Frequency (monthly)
-"Spark shell","app-20211019113801-0001","Not Recommended",1.0,569385.42,2581.57,3627,19894,571967,3503,28.41,"","JDBC[*]","","","","",1812,544575,677,19217,3.8,false,"Scan JDBCRelation(TBLS) [numPartitions=1];Execute CreateViewCommand;CollectLimit","",30
+"Spark shell","app-20211019113801-0001","Not Recommended",1.0,569439.65,2527.34,3627,19894,571967,3470,28.41,"","JDBC[*]","","","","",1812,544575,859,19035,3.68,false,"Scan JDBCRelation(TBLS) [numPartitions=1];Execute CreateViewCommand;CollectLimit","",30

core/src/test/resources/QualificationExpectations/nds_q86_fail_test_expectation.csv

@@ -1,2 +1,2 @@
 App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Estimated GPU Time Saved,SQL DF Duration,SQL Dataframe Task Duration,App Duration,GPU Opportunity,Executor CPU Time Percent,SQL Ids with Failures,Unsupported Read File Formats and Types,Unsupported Write Data Format,Complex Types,Nested Complex Types,Potential Problems,Longest SQL Duration,NONSQL Task Duration Plus Overhead,Unsupported Task Duration,Supported SQL DF Task Duration,Task Speedup Factor,App Duration Estimated,Unsupported Execs,Unsupported Expressions,Estimated Job Frequency (monthly)
-"TPC-DS Like Bench q86","app-20210319163812-1778","Not Applicable",1.36,19120.15,7050.84,9569,4320658,26171,9569,0.0,"24","","","","","",9565,3595714,0,4320658,3.8,false,"Execute CreateViewCommand","",30
+"TPC-DS Like Bench q86","app-20210319163812-1778","Not Applicable",1.36,19216.48,6954.51,9569,4320658,26171,9569,0.0,"24","","","","","",9565,3595714,0,4320658,3.66,false,"Execute CreateViewCommand","",30

core/src/test/resources/QualificationExpectations/nds_q86_fail_test_expectation_persql.csv

@@ -23,4 +23,4 @@ App Name,App ID,SQL ID,SQL Description,SQL DF Duration,GPU Opportunity,Estimated
 "TPC-DS Like Bench q86","app-20210319163812-1778",17,"Register input tables",0,0,0.0,0.0,0.0,"Not Recommended"
 "TPC-DS Like Bench q86","app-20210319163812-1778",8,"Register input tables",0,0,0.0,0.0,0.0,"Not Recommended"
 "TPC-DS Like Bench q86","app-20210319163812-1778",23,"Register input tables",0,0,0.0,0.0,0.0,"Not Recommended"
-"TPC-DS Like Bench q86","app-20210319163812-1778",24,"Benchmark Run: query=q86; iteration=0",9565,9565,2517.1,3.79,7047.89,"Not Applicable"
+"TPC-DS Like Bench q86","app-20210319163812-1778",24,"Benchmark Run: query=q86; iteration=0",9565,9565,2613.38,3.66,6951.61,"Not Applicable"

core/src/test/resources/QualificationExpectations/nds_q86_test_expectation.csv

@@ -1,2 +1,2 @@
 App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Estimated GPU Time Saved,SQL DF Duration,SQL Dataframe Task Duration,App Duration,GPU Opportunity,Executor CPU Time Percent,SQL Ids with Failures,Unsupported Read File Formats and Types,Unsupported Write Data Format,Complex Types,Nested Complex Types,Potential Problems,Longest SQL Duration,NONSQL Task Duration Plus Overhead,Unsupported Task Duration,Supported SQL DF Task Duration,Task Speedup Factor,App Duration Estimated,Unsupported Execs,Unsupported Expressions,Estimated Job Frequency (monthly)
-"TPC-DS Like Bench q86","app-20210319163812-1778","Recommended",1.36,19120.15,7050.84,9569,4320658,26171,9569,35.34,"","","","","","",9565,3595714,0,4320658,3.8,false,"Execute CreateViewCommand","",30
+"TPC-DS Like Bench q86","app-20210319163812-1778","Recommended",1.36,19216.48,6954.51,9569,4320658,26171,9569,35.34,"","","","","","",9565,3595714,0,4320658,3.66,false,"Execute CreateViewCommand","",30

core/src/test/resources/QualificationExpectations/nds_q86_test_expectation_persql.csv

@@ -1,5 +1,5 @@
 App Name,App ID,SQL ID,SQL Description,SQL DF Duration,GPU Opportunity,Estimated GPU Duration,Estimated GPU Speedup,Estimated GPU Time Saved,Recommendation
-"TPC-DS Like Bench q86","app-20210319163812-1778",24,"Benchmark Run: query=q86; iteration=0",9565,9565,2517.1,3.79,7047.89,"Strongly Recommended"
+"TPC-DS Like Bench q86","app-20210319163812-1778",24,"Benchmark Run: query=q86; iteration=0",9565,9565,2613.38,3.66,6951.61,"Strongly Recommended"
 "TPC-DS Like Bench q86","app-20210319163812-1778",0,"Register input tables",2,2,2.0,1.0,0.0,"Not Recommended"
 "TPC-DS Like Bench q86","app-20210319163812-1778",21,"Register input tables",1,1,1.0,1.0,0.0,"Not Recommended"
 "TPC-DS Like Bench q86","app-20210319163812-1778",5,"Register input tables",1,1,1.0,1.0,0.0,"Not Recommended"

core/src/test/resources/QualificationExpectations/qual_test_simple_expectation.csv

@@ -2,4 +2,4 @@ App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Esti
 "Rapids Spark Profiling Tool Unit Tests","local-1622043423018","Recommended",1.92,8472.65,7846.34,12434,132257,16319,10589,37.7,"","","JSON","","","",7143,4717,19616,112641,3.86,false,"SerializeFromObject;Execute InsertIntoHadoopFsRelationCommand json;DeserializeToObject;Filter;MapElements;Scan","",1
 "Spark shell","local-1651187225439","Not Recommended",1.0,355483.43,153.56,760,180,355637,350,87.88,"","JSON[string:bigint:int]","","","","",498,343411,97,83,1.78,false,"SerializeFromObject;CollectLimit;DeserializeToObject;Scan json;Filter;MapElements","",1
 "Spark shell","local-1651188809790","Not Recommended",1.0,166199.97,15.02,911,283,166215,45,81.18,"","JSON[string:bigint:int]","","","","UDF",715,133608,269,14,1.5,false,"CollectLimit;Scan json;Project","UDF",1
-"Rapids Spark Profiling Tool Unit Tests","local-1623281204390","Not Recommended",1.0,6240.0,0.0,2032,4666,6240,0,46.27,"","JSON[string:bigint:int]","JSON","","","UDF",1209,5793,4664,2,1.0,false,"Execute InsertIntoHadoopFsRelationCommand json;LocalTableScan;Project;Scan json;Execute CreateViewCommand","UDF",1
+"Rapids Spark Profiling Tool Unit Tests","local-1623281204390","Not Recommended",1.0,6240.0,0.0,2032,4666,6240,1,46.27,"","JSON[string:bigint:int]","JSON","","","UDF",1209,5793,4662,4,1.0,false,"Execute InsertIntoHadoopFsRelationCommand json;LocalTableScan;Project;Scan json;Execute CreateViewCommand","UDF",1

core/src/test/resources/QualificationExpectations/read_dsv1_expectation.csv

@@ -1,2 +1,2 @@
 App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Estimated GPU Time Saved,SQL DF Duration,SQL Dataframe Task Duration,App Duration,GPU Opportunity,Executor CPU Time Percent,SQL Ids with Failures,Unsupported Read File Formats and Types,Unsupported Write Data Format,Complex Types,Nested Complex Types,Potential Problems,Longest SQL Duration,NONSQL Task Duration Plus Overhead,Unsupported Task Duration,Supported SQL DF Task Duration,Task Speedup Factor,App Duration Estimated,Unsupported Execs,Unsupported Expressions,Estimated Job Frequency (monthly)
-"Spark shell","local-1624371544219","Not Recommended",1.0,174691.42,601.57,6695,20421,175293,1034,72.15,"","JSON[string:double:date:int:bigint];Text[*]","JSON","","","",1859,175857,17266,3155,2.39,false,"CollectLimit;Scan json;Execute InsertIntoHadoopFsRelationCommand json;Scan text","",30
+"Spark shell","local-1624371544219","Not Recommended",1.0,173961.11,1331.88,6695,20421,175293,2290,72.15,"","JSON[string:double:date:int:bigint];Text[*]","JSON","","","",1859,175953,13403,7018,2.39,false,"CollectLimit;Scan json;Execute InsertIntoHadoopFsRelationCommand json;Scan text","",30

core/src/test/resources/QualificationExpectations/read_dsv2_expectation.csv

@@ -1,2 +1,2 @@
 App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Estimated GPU Time Saved,SQL DF Duration,SQL Dataframe Task Duration,App Duration,GPU Opportunity,Executor CPU Time Percent,SQL Ids with Failures,Unsupported Read File Formats and Types,Unsupported Write Data Format,Complex Types,Nested Complex Types,Potential Problems,Longest SQL Duration,NONSQL Task Duration Plus Overhead,Unsupported Task Duration,Supported SQL DF Task Duration,Task Speedup Factor,App Duration Estimated,Unsupported Execs,Unsupported Expressions,Estimated Job Frequency (monthly)
-"Spark shell","local-1624371906627","Not Recommended",1.0,83172.84,565.15,6760,21802,83738,971,71.3,"","Text[*];json[double]","JSON","","","",1984,82505,18668,3134,2.39,false,"BatchScan json;Execute InsertIntoHadoopFsRelationCommand json;CollectLimit;Scan text","",30
+"Spark shell","local-1624371906627","Not Recommended",1.01,82225.67,1512.32,6760,21802,83738,2412,71.3,"","Text[*];json[double]","JSON","","","",1984,82601,13987,7815,2.68,false,"BatchScan json;Execute InsertIntoHadoopFsRelationCommand json;CollectLimit;Scan text","",30

core/src/test/resources/QualificationExpectations/spark2_expectation.csv

@@ -1,2 +1,2 @@
 App Name,App ID,Recommendation,Estimated GPU Speedup,Estimated GPU Duration,Estimated GPU Time Saved,SQL DF Duration,SQL Dataframe Task Duration,App Duration,GPU Opportunity,Executor CPU Time Percent,SQL Ids with Failures,Unsupported Read File Formats and Types,Unsupported Write Data Format,Complex Types,Nested Complex Types,Potential Problems,Longest SQL Duration,NONSQL Task Duration Plus Overhead,Unsupported Task Duration,Supported SQL DF Task Duration,Task Speedup Factor,App Duration Estimated,Unsupported Execs,Unsupported Expressions,Estimated Job Frequency (monthly)
-"Spark shell","local-1634253215009","Not Recommended",1.01,46352.24,710.75,1520,359,47063,1011,67.64,"","Text[*]","","","","",1068,44935,120,239,3.36,false,"CollectLimit;Scan text","",30
+"Spark shell","local-1634253215009","Not Recommended",1.01,46516.07,546.92,1520,359,47063,842,67.64,"","Text[*]","","","","",1068,44935,160,199,2.85,false,"CollectLimit;Scan text","",30