Add physical plan to the dot file as the graph label (#2640)

* Add physical plan to the dot file as graph label Signed-off-by: Gera Shegalov <gera@apache.org> * review * test fix
NVIDIA · Jun 9, 2021 · bb6d2ab · bb6d2ab
1 parent 73be7db
commit bb6d2ab
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 11 deletions.
diff --git a/tools/src/main/scala/com/nvidia/spark/rapids/tool/profiling/CollectInformation.scala b/tools/src/main/scala/com/nvidia/spark/rapids/tool/profiling/CollectInformation.scala
@@ -124,13 +124,16 @@ class CollectInformation(apps: ArrayBuffer[ApplicationInfo],
           list += row.getLong(1) -> row.getLong(2)
         }
 
-        for ((sqlID, planInfo) <- app.sqlPlan) {
+        val sqlPlansMap = app.sqlPlan.map { case (sqlId, sparkPlanInfo) =>
+          sqlId -> ((sparkPlanInfo, app.physicalPlanDescription(sqlId)))
+        }
+        for ((sqlID,  (planInfo, physicalPlan)) <- sqlPlansMap) {
           val dotFileWriter = new ToolTextFileWriter(outputDirectory,
             s"${app.appId}-query-$sqlID.dot")
           try {
             val metrics = map.getOrElse(sqlID, Seq.empty).toMap
-            GenerateDot.generateDotGraph(
-              QueryPlanWithMetrics(planInfo, metrics), None, dotFileWriter, sqlID + ".dot")
+            GenerateDot.generateDotGraph(QueryPlanWithMetrics(planInfo, metrics),
+              physicalPlan, None, dotFileWriter, sqlID, app.appId)
           } finally {
             dotFileWriter.close()
           }

diff --git a/tools/src/main/scala/com/nvidia/spark/rapids/tool/profiling/GenerateDot.scala b/tools/src/main/scala/com/nvidia/spark/rapids/tool/profiling/GenerateDot.scala
@@ -47,14 +47,19 @@ object GenerateDot {
    *
    * @param plan First query plan and metrics
    * @param comparisonPlan Optional second query plan and metrics
-   * @param filename Filename to write dot graph to
+   * @param sqlId id of the SQL query for the dot graph
+   * @param appId Spark application Id
    */
   def generateDotGraph(
-      plan: QueryPlanWithMetrics,
-      comparisonPlan: Option[QueryPlanWithMetrics],
-      fileWriter: ToolTextFileWriter,
-      filename: String): Unit = {
-
+    plan: QueryPlanWithMetrics,
+    physicalPlanString: String,
+    comparisonPlan: Option[QueryPlanWithMetrics],
+    fileWriter: ToolTextFileWriter,
+    sqlId: Long,
+    appId: String
+  ): Unit = {
+
+    val fileName  = s"$sqlId.dot"
     var nextId = 1
 
     def isGpuPlan(plan: SparkPlanInfo): Boolean = {
@@ -166,8 +171,25 @@ object GenerateDot {
       }
     }
 
+    val leftAlignedLabel =
+      s"""
+         |Application: $appId
+         |Query: $sqlId
+         |
+         |$physicalPlanString"""
+          .stripMargin
+          .replace("\n", "\\l")
+
     // write the dot graph to a file
-    fileWriter.write("digraph G {\n")
+    fileWriter.write(
+      s"""digraph G {
+         |
+         |label="$leftAlignedLabel"
+         |labelloc=b
+         |fontname=Courier
+         |
+         |""".stripMargin)
+
     writeGraph(fileWriter, plan, comparisonPlan.getOrElse(plan), 0)
     fileWriter.write("}\n")
   }

diff --git a/tools/src/test/scala/com/nvidia/spark/rapids/tool/profiling/GenerateDotSuite.scala b/tools/src/test/scala/com/nvidia/spark/rapids/tool/profiling/GenerateDotSuite.scala
@@ -75,7 +75,8 @@ class GenerateDotSuite extends FunSuite with BeforeAndAfterAll with Logging {
             source.close()
           }
         }
-        assert(hashAggCount === 2)
+        // 2 node labels + 1 graph label
+        assert(hashAggCount === 3)
       }
     }
   }