Stages nodes

build.sbt

@@ -14,6 +14,8 @@ spAppendScalaVersion := true
 
 libraryDependencies += "org.apache.spark" %% "spark-core" % sparkVersion.value % "provided"
 
+libraryDependencies += "org.apache.spark" %% "spark-sql" % sparkVersion.value % "provided"
+
 libraryDependencies +=  "org.apache.hadoop" % "hadoop-client" % "2.6.5" % "provided"
 
 testOptions in Test += Tests.Argument("-oF")

src/main/scala/com/qubole/sparklens/QuboleJobListener.scala

@@ -21,11 +21,13 @@ import java.net.URI
 
 import com.qubole.sparklens.analyzer._
 import com.qubole.sparklens.common.{AggregateMetrics, AppContext, ApplicationInfo}
+import com.qubole.sparklens.pluggable.ComplimentaryMetrics
 import com.qubole.sparklens.timespan.{ExecutorTimeSpan, HostTimeSpan, JobTimeSpan, StageTimeSpan}
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, Path}
 import org.apache.spark.SparkConf
 import org.apache.spark.scheduler._
+import org.apache.spark.sql.{QuboleSQLListener, SparkSession}
 
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
@@ -37,17 +39,20 @@ import scala.collection.mutable.ListBuffer
   *
   */
 
-class QuboleJobListener(sparkConf: SparkConf)  extends SparkListener {
+class QuboleJobListener(sparkConf: SparkConf) extends SparkListener {
 
-  protected val appInfo          = new ApplicationInfo()
-  protected val executorMap      = new mutable.HashMap[String, ExecutorTimeSpan]()
-  protected val hostMap          = new mutable.HashMap[String, HostTimeSpan]()
-  protected val jobMap           = new mutable.HashMap[Long, JobTimeSpan]
-  protected val jobSQLExecIDMap  = new mutable.HashMap[Long, Long]
-  protected val stageMap         = new mutable.HashMap[Int, StageTimeSpan]
-  protected val stageIDToJobID   = new mutable.HashMap[Int, Long]
-  protected val failedStages     = new ListBuffer[String]
-  protected val appMetrics       = new AggregateMetrics()
+  protected val appInfo               = new ApplicationInfo()
+  protected val executorMap           = new mutable.HashMap[String, ExecutorTimeSpan]()
+  protected val hostMap               = new mutable.HashMap[String, HostTimeSpan]()
+  protected val jobMap                = new mutable.HashMap[Long, JobTimeSpan]
+  protected val jobSQLExecIDMap       = new mutable.HashMap[Long, Long]
+  protected val stageMap              = new mutable.HashMap[Int, StageTimeSpan]
+  protected val stageIDToJobID        = new mutable.HashMap[Int, Long]
+  protected val failedStages          = new ListBuffer[String]
+  protected val appMetrics            = new AggregateMetrics()
+  val pluggableMetricsMap   = new mutable.HashMap[String, ComplimentaryMetrics]()
+
+  private var sqlListener: QuboleSQLListener = _
 
   private def hostCount():Int = hostMap.size
 
@@ -128,6 +133,7 @@ class QuboleJobListener(sparkConf: SparkConf)  extends SparkListener {
     if (taskEnd.taskInfo.failed) {
       //println(s"\nTask Failed \n ${taskEnd.reason}"
     }
+    sqlListener.onTaskEnd(taskEnd)
   }
 
   private[this] def dumpData(appContext: AppContext): Unit = {
@@ -144,6 +150,11 @@ class QuboleJobListener(sparkConf: SparkConf)  extends SparkListener {
     //println(s"Application ${applicationStart.appId} started at ${applicationStart.time}")
     appInfo.applicationID = applicationStart.appId.getOrElse("NA")
     appInfo.startTime     = applicationStart.time
+
+    val spark = SparkSession.builder().getOrCreate()
+    sqlListener = new QuboleSQLListener(sparkConf, this)
+    spark.listenerManager.register(sqlListener)
+    sqlListener.onApplicationStart(applicationStart)
   }
 
   override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
@@ -157,7 +168,8 @@ class QuboleJobListener(sparkConf: SparkConf)  extends SparkListener {
       jobMap,
       jobSQLExecIDMap,
       stageMap,
-      stageIDToJobID)
+      stageIDToJobID,
+      pluggableMetricsMap)
 
     asyncReportingEnabled(sparkConf) match {
       case true => {
@@ -245,5 +257,6 @@ class QuboleJobListener(sparkConf: SparkConf)  extends SparkListener {
       jobTimeSpan.addStage(stageTimeSpan)
       stageTimeSpan.finalUpdate()
     }
+    sqlListener.onStageCompleted(stageCompleted)
   }
 }

src/main/scala/com/qubole/sparklens/QuboleNotebookListener.scala

@@ -125,7 +125,8 @@ class QuboleNotebookListener(sparkConf: SparkConf) extends QuboleJobListener(spa
       jobMap,
       jobSQLExecIDMap,
       stageMap,
-      stageIDToJobID)
+      stageIDToJobID,
+      pluggableMetricsMap)
 
     val out = new mutable.StringBuilder()
 

src/main/scala/com/qubole/sparklens/common/AppContext.scala

@@ -16,6 +16,7 @@
 */
 package com.qubole.sparklens.common
 
+import com.qubole.sparklens.pluggable.ComplimentaryMetrics
 import com.qubole.sparklens.timespan._
 import org.json4s.DefaultFormats
 import org.json4s.JsonAST.JValue
@@ -24,14 +25,15 @@ import org.json4s.MappingException
 
 import scala.collection.mutable
 
-case class AppContext(appInfo:        ApplicationInfo,
-                      appMetrics:     AggregateMetrics,
-                      hostMap:        mutable.HashMap[String, HostTimeSpan],
-                      executorMap:    mutable.HashMap[String, ExecutorTimeSpan],
-                      jobMap:         mutable.HashMap[Long, JobTimeSpan],
-                      jobSQLExecIdMap:mutable.HashMap[Long, Long],
-                      stageMap:       mutable.HashMap[Int, StageTimeSpan],
-                      stageIDToJobID: mutable.HashMap[Int, Long]) {
+case class AppContext(appInfo:              ApplicationInfo,
+                      appMetrics:           AggregateMetrics,
+                      hostMap:              mutable.HashMap[String, HostTimeSpan],
+                      executorMap:          mutable.HashMap[String, ExecutorTimeSpan],
+                      jobMap:               mutable.HashMap[Long, JobTimeSpan],
+                      jobSQLExecIdMap:      mutable.HashMap[Long, Long],
+                      stageMap:             mutable.HashMap[Int, StageTimeSpan],
+                      stageIDToJobID:       mutable.HashMap[Int, Long],
+                      pluggableMetricsMap:  mutable.HashMap[String, ComplimentaryMetrics]) {
 
   def filterByStartAndEndTime(startTime: Long, endTime: Long): AppContext = {
     new AppContext(appInfo,
@@ -48,7 +50,8 @@ case class AppContext(appInfo:        ApplicationInfo,
       stageMap
         .filter(x => x._2.startTime >= startTime &&
                      x._2.endTime <= endTime),
-      stageIDToJobID)
+      stageIDToJobID,
+      pluggableMetricsMap)
   }
 
   override def toString(): String = {
@@ -61,7 +64,8 @@ case class AppContext(appInfo:        ApplicationInfo,
       "jobMap" -> AppContext.getMap(jobMap),
       "jobSQLExecIdMap" -> jobSQLExecIdMap,
       "stageMap" -> AppContext.getMap(stageMap),
-      "stageIDToJobID" -> stageIDToJobID
+      "stageIDToJobID" -> stageIDToJobID,
+      "pluggableMetricsMap" -> ComplimentaryMetrics.getMap(pluggableMetricsMap)
     )
     Serialization.writePretty(map)
   }
@@ -134,7 +138,8 @@ object AppContext {
       JobTimeSpan.getTimeSpan((json \ "jobMap").extract[Map[String, JValue]]),
       getJobSQLExecIdMap(json, new mutable.HashMap[Long, Long]),
       StageTimeSpan.getTimeSpan((json \ "stageMap").extract[Map[String, JValue]]),
-      getJobToStageMap((json \ "stageIDToJobID").extract[Map[Int, JValue]])
+      getJobToStageMap((json \ "stageIDToJobID").extract[Map[Int, JValue]]),
+      ComplimentaryMetrics.getMetricsMap((json \ "pluggableMetricsMap").extract[Map[String, JValue]])
     )
 }
 

src/main/scala/com/qubole/sparklens/pluggable/ComplimentaryMetrics.scala

@@ -0,0 +1,57 @@
+package com.qubole.sparklens.pluggable
+
+import org.json4s.{DefaultFormats, MappingException}
+import org.json4s.JsonAST.JValue
+
+import scala.collection.mutable
+
+trait ComplimentaryMetrics {
+  def getMap(): Map[String, _ <: Any] = {
+    throw new NotImplementedError(s"getMap() method is not implemented.")
+  }
+
+  def getObject(json: JValue): ComplimentaryMetrics = {
+    throw new NotImplementedError(s"getObject() method is not implemented.")
+  }
+}
+
+object ComplimentaryMetrics {
+
+  /**
+    * Returns object which extends [[ComplimentaryMetrics]] by matching input string
+    */
+  def fromString(value: String): ComplimentaryMetrics = {
+    value.toLowerCase match {
+      case "sqlmetrics" => SQLMetrics
+      case _ => throw new Exception(s"Object ${value} not found.")
+    }
+  }
+
+  /**
+    * Used for for extracting the pluggableMetricsMap in [[com.qubole.sparklens.QuboleJobListener]]
+    * to construct [[com.qubole.sparklens.common.AppContext]] from the JSON.
+    */
+  def getMetricsMap(json: Map[String, JValue]): mutable.HashMap[String, ComplimentaryMetrics] = {
+    val metricsMap = new mutable.HashMap[String, ComplimentaryMetrics]
+    try {
+      implicit val formats = DefaultFormats
+      val metricsMap = new mutable.HashMap[String, ComplimentaryMetrics]
+      json.keys.map(key => {
+        val value = json.get(key).get
+        metricsMap.put(key, fromString(key).getObject(value))
+      })
+    } catch {
+      case e: Exception if !e.isInstanceOf[MappingException] =>
+        throw(e)
+    }
+    metricsMap
+  }
+
+  /**
+    * Used for for converting the pluggableMetricsMap in [[com.qubole.sparklens.QuboleJobListener]]
+    * to a formatted map which is then dumped in the JSON file/printed on console.
+    */
+  def getMap(metricsMap: mutable.HashMap[String, _ <: ComplimentaryMetrics]): Map[String, Any] = {
+    metricsMap.keys.map(key => (key.toString, metricsMap(key).getMap)).toMap
+  }
+}