Embed the viz in the UI in a toggleable manner

core/src/main/resources/org/apache/spark/ui/static/spark-stage-viz.js

@@ -15,6 +15,31 @@
  * limitations under the License.
  */
 
+var stageVizIsRendered = false
+
+/*
+ * Render or remove the stage visualization on the UI.
+ * This assumes that the visualization is stored in the "#viz-graph" element.
+ */
+function toggleStageViz() {
+  $(".expand-visualization-arrow").toggleClass('arrow-closed');
+  $(".expand-visualization-arrow").toggleClass('arrow-open');
+  var shouldRender = $(".expand-visualization-arrow").hasClass("arrow-open");
+  if (shouldRender) {
+    // If the viz is already rendered, just show it
+    if (stageVizIsRendered) {
+      $("#viz-graph").show();
+    } else {
+      renderStageViz();
+      stageVizIsRendered = true;
+    }
+  } else {
+    // Instead of emptying the element once and for all, cache it for use
+    // again later in case we want to expand the visualization again
+    $("#viz-graph").hide();
+  }
+}
+
 /*
  * Render a DAG that describes the RDDs for a given stage.
  *
@@ -27,6 +52,14 @@
  */
 function renderStageViz() {
 
+  // If there is not a dot file to render, report error
+  if (d3.select("#viz-dot-file").empty()) {
+    d3.select("#viz-graph")
+      .append("div")
+      .text("No visualization information available for this stage.");
+    return;
+  }
+
   // Parse the dot file and render it in an SVG
   var dot = d3.select("#viz-dot-file").text();
   var escaped_dot = dot
@@ -35,7 +68,7 @@ function renderStageViz() {
     .replace(/"/g, "\"");
   var g = graphlibDot.read(escaped_dot);
   var render = new dagreD3.render();
-  var svg = d3.select("#viz-graph");
+  var svg = d3.select("#viz-graph").append("svg");
   svg.call(render, g);
 
   // Set the appropriate SVG dimensions to ensure that all elements are displayed
@@ -88,7 +121,6 @@ function renderStageViz() {
   var endY = toFloat(svg.style("height")) + svgMargin;
   var newViewBox = startX + " " + startY + " " + endX + " " + endY;
   svg.attr("viewBox", newViewBox);
-
 }
 
 /* Helper method to convert attributes to numeric values. */

core/src/main/resources/org/apache/spark/ui/static/webui.css

@@ -145,7 +145,7 @@ pre {
   border: none;
 }
 
-span.expand-additional-metrics {
+span.expand-additional-metrics, span.expand-visualization {
   cursor: pointer;
 }
 

core/src/main/scala/org/apache/spark/ui/SparkUI.scala

@@ -147,7 +147,7 @@ private[spark] object SparkUI {
     val storageStatusListener = new StorageStatusListener
     val executorsListener = new ExecutorsListener(storageStatusListener)
     val storageListener = new StorageListener(storageStatusListener)
-    val visualizationListener = new VisualizationListener
+    val visualizationListener = new VisualizationListener(conf)
 
     listenerBus.addListener(environmentListener)
     listenerBus.addListener(storageStatusListener)

core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala

@@ -37,20 +37,29 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
   private val vizListener = parent.vizListener
 
   /**
-   * Return a DOM element that contains an RDD DAG visualization for this stage.
-   * If there is no visualization information for this stage, return an empty element.
+   * Return a DOM element containing the "Show Visualization" toggle that, if enabled,
+   * renders the visualization for the stage. If there is no visualization information
+   * available for this stage, an appropriate message is displayed to the user.
    */
-  private def renderViz(stageId: Int): Seq[Node] = {
+  private def showVizElement(stageId: Int): Seq[Node] = {
     val graph = vizListener.getVizGraph(stageId)
-    if (graph.isEmpty) {
-      Seq.empty
-    } else {
-      <div id="viz-dot-file" style="display:none">
-        {VizGraph.makeDotFile(graph.get)}
-      </div>
-      <svg id="viz-graph"></svg>
-      <script type="text/javascript">renderStageViz()</script>
-    }
+    <div>
+      <span class="expand-visualization" onclick="render();">
+        <span class="expand-visualization-arrow arrow-closed"></span>
+        <strong>Show Visualization</strong>
+      </span>
+      <div id="viz-graph"></div>
+      <script type="text/javascript">
+        {Unparsed(s"function render() { toggleStageViz(); }")}
+      </script>
+      {
+        if (graph.isDefined) {
+          <div id="viz-dot-file" style="display:none">
+            {VizGraph.makeDotFile(graph.get)}
+          </div>
+        }
+      }
+    </div>
   }
 
   def render(request: HttpServletRequest): Seq[Node] = {
@@ -75,8 +84,6 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
           s"Details for Stage $stageId (Attempt $stageAttemptId)", content, parent)
       }
 
-      val viz: Seq[Node] = renderViz(stageId)
-
       val stageData = stageDataOption.get
       val tasks = stageData.taskData.values.toSeq.sortBy(_.taskInfo.launchTime)
 
@@ -453,9 +460,9 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
         if (accumulables.size > 0) { <h4>Accumulators</h4> ++ accumulableTable } else Seq()
 
       val content =
-        viz ++
         summary ++
         showAdditionalMetrics ++
+        showVizElement(stageId) ++
         <h4>Summary Metrics for {numCompleted} Completed Tasks</h4> ++
         <div>{summaryTable.getOrElse("No tasks have reported metrics yet.")}</div> ++
         <h4>Aggregated Metrics by Executor</h4> ++ executorTable.toNodeSeq ++

core/src/main/scala/org/apache/spark/ui/viz/VisualizationListener.scala

@@ -19,14 +19,15 @@ package org.apache.spark.ui.viz
 
 import scala.collection.mutable
 
+import org.apache.spark.SparkConf
 import org.apache.spark.scheduler._
 import org.apache.spark.ui.SparkUI
 
 /**
  * A SparkListener that constructs a graph of the RDD DAG for each stage.
  * This graph will be used for rendering visualization in the UI later.
  */
-private[ui] class VisualizationListener extends SparkListener {
+private[ui] class VisualizationListener(conf: SparkConf) extends SparkListener {
 
   // A list of stage IDs to track the order in which stages are inserted
   private val stageIds = new mutable.ArrayBuffer[Int]

core/src/main/scala/org/apache/spark/ui/viz/VizGraph.scala

@@ -71,7 +71,7 @@ private[ui] object VizGraph {
 
     // Populate nodes, edges, and scopes
     rddInfos.foreach { rdd =>
-      val node = nodes.getOrElseUpdate(rdd.id, VizNode(rdd.id, rdd.name, rdd.isCached))
+      val node = nodes.getOrElseUpdate(rdd.id, VizNode(rdd.id, rdd.name))
       edges ++= rdd.parentIds.map { parentId => VizEdge(parentId, rdd.id) }
 
       if (rdd.scope == null) {