Qualification tool: Enhance mapping of Execs to stages

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

@@ -198,6 +198,78 @@ class QualificationAppInfo(
     }
   }
 
+  private def singleExecMatch(execInfo: ExecInfo): (Seq[(Int, ExecInfo)], Option[ExecInfo]) = {
+    val associatedStages = {
+      if (execInfo.stages.size > 1) {
+        execInfo.stages.toSeq
+      } 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.")
+        Seq.empty
+      } else {
+        Seq(execInfo.stages.head)
+      }
+    }
+    if (associatedStages.nonEmpty) {
+      (associatedStages.map((_, execInfo)), None)
+    } else {
+      (Seq.empty, Some(execInfo))
+    }
+  }
+
+  private def doubleExecMatch(neighbor: ExecInfo, execInfo: ExecInfo): (
+    Seq[(Int, ExecInfo)], Option[ExecInfo]) = {
+    val associatedStages = {
+      if (execInfo.stages.size > 1) {
+        execInfo.stages.toSeq
+      } else if (execInfo.stages.size < 1) {
+        if (neighbor.stages.size >= 1) {
+          neighbor.stages.headOption.toSeq
+        } 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.")
+          Seq.empty
+        }
+      } else {
+        Seq(execInfo.stages.head)
+      }
+    }
+    if (associatedStages.nonEmpty) {
+      (associatedStages.map((_, execInfo)), None)
+    } else {
+      (Seq.empty, Some(execInfo))
+    }
+  }
+
+  private def tripleExecMatch(prev: ExecInfo, execInfo: ExecInfo, next: ExecInfo):
+  (Seq[(Int, ExecInfo)], Option[ExecInfo]) = {
+    val associatedStages = {
+      if (execInfo.stages.size > 1) {
+        execInfo.stages.toSeq
+      } else if (execInfo.stages.size < 1) {
+        if (prev.stages.size >= 1) {
+          prev.stages.headOption.toSeq
+        } else if (next.stages.size >= 1) {
+          next.stages.headOption.toSeq
+        } 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.")
+          Seq.empty
+        }
+      } else {
+        Seq(execInfo.stages.head)
+      }
+    }
+    if (associatedStages.nonEmpty) {
+      (associatedStages.map((_, execInfo)), None)
+    } else {
+      (Seq.empty, Some(execInfo))
+    }
+  }
+
   private def getStageToExec(execInfos: Seq[ExecInfo]): (Map[Int, Seq[ExecInfo]], Seq[ExecInfo]) = {
     val execsWithoutStages = new ArrayBuffer[ExecInfo]()
 
@@ -209,69 +281,34 @@ class QualificationAppInfo(
     // 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))
+    val execsToStageMap = execInfosInOrder.indices.map {
+      // corner case to handle first element
+      case 0 => if (execInfosInOrder.size > 1) {
+        // If there are more than one Execs, then check if the next Exec has a stageId.
+        doubleExecMatch(execInfosInOrder(1), execInfosInOrder(0))
+      } else {
+        singleExecMatch(execInfosInOrder(0))
+      }
+      // corner case to handle last element
+      case i if i == execInfosInOrder.size - 1 =>
+        if (execInfosInOrder.size > 1) {
+          // If there are more than one Execs, then check if the previous Exec has a stageId.
+          doubleExecMatch(execInfosInOrder(i - 1), execInfosInOrder(i))
+        } else {
+          singleExecMatch(execInfosInOrder(i))
+        }
+      case i =>
+        tripleExecMatch(execInfosInOrder(i - 1), execInfosInOrder(i), execInfosInOrder(i + 1))
+      case _ =>
+        (Seq.empty, None)
     }
+    val perStageSum = execsToStageMap.map(_._1).toList.flatten
+      .groupBy(_._1).map { case (stage, execInfo) =>
+      (stage, execInfo.map(_._2))}
+
+    // Add all the execs that don't have a stageId to execsWithoutStages.
+    execsWithoutStages ++=  execsToStageMap.map(_._2).toList.flatten
+
     (perStageSum, execsWithoutStages)
   }
 

core/src/test/resources/QualificationExpectations/complex_dec_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-1626104300434","Not Recommended",1.01,129484.66,1619.33,2429,1469,131104,2429,88.35,"","","","struct<firstname:string,middlename:array<string>,lastname:string>;struct<current:struct<state:string,city:string>,previous:struct<state:map<string,string>,city:string>>;array<struct<city:string,state:string>>;map<string,string>;map<string,array<string>>;map<string,map<string,string>>;array<array<string>>;array<string>","struct<firstname:string,middlename:array<string>,lastname:string>;struct<current:struct<state:string,city:string>,previous:struct<state:map<string,string>,city:string>>;array<struct<city:string,state:string>>;map<string,array<string>>;map<string,map<string,string>>;array<array<string>>","NESTED COMPLEX TYPE",1260,128847,0,1469,3.0,false,"CollectLimit","",30
+"Spark shell","local-1626104300434","Not Recommended",1.0,129898.52,1205.47,2429,1469,131104,1923,88.35,"","","","struct<firstname:string,middlename:array<string>,lastname:string>;struct<current:struct<state:string,city:string>,previous:struct<state:map<string,string>,city:string>>;array<struct<city:string,state:string>>;map<string,string>;map<string,array<string>>;map<string,map<string,string>>;array<array<string>>;array<string>","struct<firstname:string,middlename:array<string>,lastname:string>;struct<current:struct<state:string,city:string>,previous:struct<state:map<string,string>,city:string>>;array<struct<city:string,state:string>>;map<string,array<string>>;map<string,map<string,string>>;array<array<string>>","NESTED COMPLEX TYPE",1260,128847,306,1163,2.68,false,"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,19216.48,6954.51,9569,4320658,26171,9569,0.0,"24","","","","","",9565,3595714,0,4320658,3.66,false,"Execute CreateViewCommand","",30
+"TPC-DS Like Bench q86","app-20210319163812-1778","Not Applicable",1.36,19230.84,6940.15,9569,4320658,26171,9569,0.0,"24","","","","","",9565,3595714,0,4320658,3.64,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,2613.38,3.66,6951.61,"Not Applicable"
+"TPC-DS Like Bench q86","app-20210319163812-1778",24,"Benchmark Run: query=q86; iteration=0",9565,9565,2627.74,3.64,6937.25,"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,19216.48,6954.51,9569,4320658,26171,9569,35.34,"","","","","","",9565,3595714,0,4320658,3.66,false,"Execute CreateViewCommand","",30
+"TPC-DS Like Bench q86","app-20210319163812-1778","Recommended",1.36,19230.84,6940.15,9569,4320658,26171,9569,35.34,"","","","","","",9565,3595714,0,4320658,3.64,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,2613.38,3.66,6951.61,"Strongly Recommended"
+"TPC-DS Like Bench q86","app-20210319163812-1778",24,"Benchmark Run: query=q86; iteration=0",9565,9565,2627.74,3.64,6937.25,"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

@@ -1,5 +1,5 @@
 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)
-"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,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
+"Rapids Spark Profiling Tool Unit Tests","local-1622043423018","Recommended",1.92,8488.68,7830.31,12434,132257,16319,10577,37.7,"","","JSON","","","",7143,4717,19744,112513,3.85,false,"SerializeFromObject;Execute InsertIntoHadoopFsRelationCommand json;DeserializeToObject;Filter;MapElements;Scan","",1
+"Spark shell","local-1651187225439","Not Recommended",1.0,355550.33,86.66,760,180,355637,253,87.88,"","JSON[string:bigint:int]","","","","",498,343411,120,60,1.52,false,"SerializeFromObject;CollectLimit;DeserializeToObject;Scan json;Filter;MapElements","",1
+"Spark shell","local-1651188809790","Not Recommended",1.0,166205.19,9.8,911,283,166215,38,81.18,"","JSON[string:bigint:int]","","","","UDF",715,133608,271,12,1.34,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,2,46.27,"","JSON[string:bigint:int]","JSON","","","UDF",1209,5793,4661,5,1.0,false,"Execute InsertIntoHadoopFsRelationCommand json;LocalTableScan;Project;Scan json;Execute CreateViewCommand","UDF",1

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

@@ -2,12 +2,12 @@ App Name,App ID,SQL ID,SQL Description,SQL DF Duration,GPU Opportunity,Estimated
 "Rapids Spark Profiling Tool Unit Tests","local-1622043423018",1,"count at QualificationInfoUtils.scala:94",7143,6719,2078.49,3.43,5064.5,"Strongly Recommended"
 "Rapids Spark Profiling Tool Unit Tests","local-1622043423018",3,"count at QualificationInfoUtils.scala:94",2052,1660,800.56,2.56,1251.43,"Strongly Recommended"
 "Rapids Spark Profiling Tool Unit Tests","local-1622043423018",2,"count at QualificationInfoUtils.scala:94",1933,1551,763.96,2.53,1169.03,"Strongly Recommended"
-"Spark shell","local-1651187225439",0,"show at <console>:26",498,249,373.5,1.33,124.5,"Recommended"
-"Spark shell","local-1651188809790",1,"show at <console>:26",196,98,147.0,1.33,49.0,"Recommended"
-"Spark shell","local-1651187225439",1,"show at <console>:26",262,60,240.54,1.08,21.45,"Not Recommended"
-"Rapids Spark Profiling Tool Unit Tests","local-1622043423018",0,"json at QualificationInfoUtils.scala:76",1306,187,1246.97,1.04,59.02,"Not Recommended"
+"Spark shell","local-1651188809790",1,"show at <console>:26",196,75,165.75,1.18,30.24,"Not Recommended"
+"Spark shell","local-1651187225439",0,"show at <console>:26",498,168,430.53,1.15,67.46,"Not Recommended"
+"Spark shell","local-1651187225439",1,"show at <console>:26",262,80,240.22,1.09,21.77,"Not Recommended"
+"Rapids Spark Profiling Tool Unit Tests","local-1622043423018",0,"json at QualificationInfoUtils.scala:76",1306,164,1267.14,1.03,38.85,"Not Recommended"
 "Rapids Spark Profiling Tool Unit Tests","local-1623281204390",0,"json at QualificationInfoUtils.scala:130",1209,0,1209.0,1.0,0.0,"Not Recommended"
-"Spark shell","local-1651188809790",0,"show at <console>:26",715,2,715.0,1.0,0.0,"Not Recommended"
+"Spark shell","local-1651188809790",0,"show at <console>:26",715,5,715.0,1.0,0.0,"Not Recommended"
 "Rapids Spark Profiling Tool Unit Tests","local-1623281204390",2,"json at QualificationInfoUtils.scala:136",321,0,321.0,1.0,0.0,"Not Recommended"
 "Rapids Spark Profiling Tool Unit Tests","local-1623281204390",5,"json at QualificationInfoUtils.scala:136",129,0,129.0,1.0,0.0,"Not Recommended"
 "Rapids Spark Profiling Tool Unit Tests","local-1623281204390",8,"json at QualificationInfoUtils.scala:136",127,0,127.0,1.0,0.0,"Not Recommended"

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,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
+"Spark shell","local-1624371544219","Not Recommended",1.0,174006.51,1286.48,6695,20421,175293,2268,72.15,"","JSON[string:double:date:int:bigint];Text[*]","JSON","","","",1859,175953,13469,6952,2.31,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.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
+"Spark shell","local-1624371906627","Not Recommended",1.01,82304.74,1433.25,6760,21802,83738,2388,71.3,"","Text[*];json[double]","JSON","","","",1984,82601,14064,7738,2.5,false,"BatchScan json;Execute InsertIntoHadoopFsRelationCommand json;CollectLimit;Scan text","",30