Qualification tool: update SQL Df value used and look at jobs in SQL

tools/src/main/scala/com/nvidia/spark/rapids/tool/planparser/WholeStageExecParser.scala

@@ -39,17 +39,20 @@ case class WholeStageExecParser(
     val maxDuration = SQLPlanParser.getTotalDuration(accumId, app)
     val stagesInNode = SQLPlanParser.getStagesInSQLNode(node, app)
 
-    val childNodeRes = node.nodes.flatMap { c =>
+    val childNodes = node.nodes.flatMap { c =>
       SQLPlanParser.parsePlanNode(c, sqlID, checker, app)
     }
     // if any of the execs in WholeStageCodegen supported mark this entire thing
     // as supported
-    val anySupported = childNodeRes.exists(_.isSupported == true)
+    val anySupported = childNodes.exists(_.isSupported == true)
     // average speedup across the execs in the WholeStageCodegen for now
-    val supportedChildren = childNodeRes.filterNot(_.shouldRemove)
+    val supportedChildren = childNodes.filterNot(_.shouldRemove)
     val avSpeedupFactor = SQLPlanParser.averageSpeedup(supportedChildren.map(_.speedupFactor))
+    // can't rely on the wholeStagecodeGen having a stage if children do so aggregate them together
+    // for now
+    val allStagesIncludingChildren = childNodes.flatMap(_.stages).toSet ++ stagesInNode.toSet
     val execInfo = new ExecInfo(sqlID, node.name, node.name, avSpeedupFactor, maxDuration,
-      node.id, anySupported, Some(childNodeRes), stagesInNode)
+      node.id, anySupported, Some(childNodes), allStagesIncludingChildren.toSeq)
     Seq(execInfo)
   }
 }

tools/src/main/scala/com/nvidia/spark/rapids/tool/qualification/PluginTypeChecker.scala

@@ -200,12 +200,7 @@ class PluginTypeChecker extends Logging {
         // check if any of the not supported types are in the schema
         val nsFiltered = dtSupMap(NS).filter(t => schemaLower.contains(t.toLowerCase()))
         if (nsFiltered.nonEmpty) {
-          val deDuped = if (nsFiltered.contains("dec") && nsFiltered.contains("decimal")) {
-            nsFiltered.filterNot(_.equals("dec"))
-          } else {
-            nsFiltered
-          }
-          (0.0, deDuped.toSet)
+          (0.0, nsFiltered.toSet)
         } else {
           // Started out giving different weights based on partial support and so forth
           // but decided to be optimistic and not penalize if we don't know, perhaps

tools/src/main/scala/com/nvidia/spark/rapids/tool/qualification/QualOutputWriter.scala

@@ -75,7 +75,7 @@ class QualOutputWriter(outputDir: String, reportReadSchema: Boolean, printStdout
       val plans = sums.flatMap(_.planInfo)
       val allExecs = QualOutputWriter.getAllExecsFromPlan(plans)
       val headersAndSizes = QualOutputWriter
-        .getDetailedExecsHeaderStringsAndSizes(sums, allExecs)
+        .getDetailedExecsHeaderStringsAndSizes(sums, allExecs.toSeq)
       csvFileWriter.write(QualOutputWriter.constructDetailedHeader(headersAndSizes, ",", false))
       sums.foreach { sumInfo =>
         val rows = QualOutputWriter.constructExecsInfo(sumInfo, headersAndSizes, ",", false)
@@ -165,6 +165,7 @@ object QualOutputWriter {
   val ESTIMATED_GPU_DURATION = "Estimated GPU Duration"
   val ESTIMATED_GPU_SPEEDUP = "Estimated GPU Speedup"
   val ESTIMATED_GPU_TIMESAVED = "Estimated GPU Time Saved"
+  val STAGE_ESTIMATED_STR = "Stage Estimated"
 
   val APP_DUR_STR_SIZE: Int = APP_DUR_STR.size
   val SQL_DUR_STR_SIZE: Int = SQL_DUR_STR.size
@@ -360,7 +361,8 @@ object QualOutputWriter {
       STAGE_ID_STR -> STAGE_ID_STR.size,
       AVERAGE_SPEEDUP_STR -> AVERAGE_SPEEDUP_STR.size,
       STAGE_DUR_STR -> STAGE_DUR_STR.size,
-      UNSUPPORTED_TASK_DURATION_STR -> UNSUPPORTED_TASK_DURATION_STR.size
+      UNSUPPORTED_TASK_DURATION_STR -> UNSUPPORTED_TASK_DURATION_STR.size,
+      STAGE_ESTIMATED_STR -> STAGE_ESTIMATED_STR.size
     )
     detailedHeadersAndFields
   }
@@ -377,23 +379,24 @@ object QualOutputWriter {
         info.stageId.toString -> headersAndSizes(STAGE_ID_STR),
         f"${info.averageSpeedup}%1.2f"  -> headersAndSizes(AVERAGE_SPEEDUP_STR),
         info.stageTaskTime.toString -> headersAndSizes(STAGE_DUR_STR),
-        info.unsupportedTaskDur.toString -> headersAndSizes(UNSUPPORTED_TASK_DURATION_STR))
+        info.unsupportedTaskDur.toString -> headersAndSizes(UNSUPPORTED_TASK_DURATION_STR),
+        info.estimated.toString -> headersAndSizes(STAGE_ESTIMATED_STR))
       constructOutputRow(data, delimiter, prettyPrint)
     }
   }
 
-  def getAllExecsFromPlan(plans: Seq[PlanInfo]): Seq[ExecInfo] = {
+  def getAllExecsFromPlan(plans: Seq[PlanInfo]): Set[ExecInfo] = {
     val topExecInfo = plans.flatMap(_.execInfo)
     topExecInfo.flatMap { e =>
       e.children.getOrElse(Seq.empty) :+ e
-    }
+    }.toSet
   }
 
   def constructExecsInfo(
       sumInfo: QualificationSummaryInfo,
       headersAndSizes: LinkedHashMap[String, Int],
       delimiter: String = "|",
-      prettyPrint: Boolean): Seq[String] = {
+      prettyPrint: Boolean): Set[String] = {
     val allExecs = getAllExecsFromPlan(sumInfo.planInfo)
     val appId = sumInfo.appId
     allExecs.flatMap { info =>

tools/src/main/scala/org/apache/spark/sql/rapids/tool/AppBase.scala

@@ -47,6 +47,7 @@ abstract class AppBase(
 
   // jobId to job info
   val jobIdToInfo = new HashMap[Int, JobInfoClass]()
+  val jobIdToSqlID: HashMap[Int, Long] = HashMap.empty[Int, Long]
 
   // SQL containing any Dataset operation or RDD to DataSet/DataFrame operation
   val sqlIDToDataSetOrRDDCase: HashSet[Long] = HashSet[Long]()

tools/src/main/scala/org/apache/spark/sql/rapids/tool/EventProcessorBase.scala

@@ -318,6 +318,7 @@ abstract class EventProcessorBase[T <: AppBase](app: T) extends SparkListener wi
       ProfileUtils.isPluginEnabled(event.properties.asScala) || app.gpuMode
     )
     app.jobIdToInfo.put(event.jobId, thisJob)
+    sqlID.foreach(app.jobIdToSqlID(event.jobId) = _)
   }
 
   override def onJobStart(jobStart: SparkListenerJobStart): Unit = {