Added Json output and associated tests.

Author: kayousterhout

core/src/main/scala/org/apache/spark/CacheManager.scala

@@ -112,7 +112,7 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
           logInfo(s"Whoever was loading $id failed; we'll try it ourselves")
           loading.add(id)
         }
-        values.data.map(_.asInstanceOf[Iterator[T]])
+        values.map(_.data.asInstanceOf[Iterator[T]])
       }
     }
   }

core/src/main/scala/org/apache/spark/executor/TaskMetrics.scala

@@ -109,7 +109,7 @@ private[spark] object DataReadMethod extends Enumeration with Serializable {
  * Metrics about reading input data.
  */
 @DeveloperApi
-case class InputMetrics(val readMethod: DataReadMethod.Value) {
+case class InputMetrics(readMethod: DataReadMethod.Value) {
   /**
    * Total bytes read.
    */

core/src/main/scala/org/apache/spark/util/JsonProtocol.scala

@@ -26,7 +26,8 @@ import org.json4s.DefaultFormats
 import org.json4s.JsonDSL._
 import org.json4s.JsonAST._
 
-import org.apache.spark.executor.{ShuffleReadMetrics, ShuffleWriteMetrics, TaskMetrics}
+import org.apache.spark.executor.{DataReadMethod, InputMetrics, ShuffleReadMetrics,
+  ShuffleWriteMetrics, TaskMetrics}
 import org.apache.spark.scheduler._
 import org.apache.spark.storage._
 import org.apache.spark._
@@ -213,6 +214,8 @@ private[spark] object JsonProtocol {
       taskMetrics.shuffleReadMetrics.map(shuffleReadMetricsToJson).getOrElse(JNothing)
     val shuffleWriteMetrics =
       taskMetrics.shuffleWriteMetrics.map(shuffleWriteMetricsToJson).getOrElse(JNothing)
+    val inputMetrics =
+      taskMetrics.inputMetrics.map(inputMetricsToJson).getOrElse(JNothing)
     val updatedBlocks =
       taskMetrics.updatedBlocks.map { blocks =>
         JArray(blocks.toList.map { case (id, status) =>
@@ -230,6 +233,7 @@ private[spark] object JsonProtocol {
     ("Disk Bytes Spilled" -> taskMetrics.diskBytesSpilled) ~
     ("Shuffle Read Metrics" -> shuffleReadMetrics) ~
     ("Shuffle Write Metrics" -> shuffleWriteMetrics) ~
+    ("Input Metrics" -> inputMetrics) ~
     ("Updated Blocks" -> updatedBlocks)
   }
 
@@ -247,6 +251,11 @@ private[spark] object JsonProtocol {
     ("Shuffle Write Time" -> shuffleWriteMetrics.shuffleWriteTime)
   }
 
+  def inputMetricsToJson(inputMetrics: InputMetrics): JValue = {
+    ("Data Read Method" -> inputMetrics.readMethod.toString) ~
+    ("Bytes Read" -> inputMetrics.bytesRead)
+  }
+
   def taskEndReasonToJson(taskEndReason: TaskEndReason): JValue = {
     val reason = Utils.getFormattedClassName(taskEndReason)
     val json = taskEndReason match {
@@ -528,6 +537,8 @@ private[spark] object JsonProtocol {
       Utils.jsonOption(json \ "Shuffle Read Metrics").map(shuffleReadMetricsFromJson)
     metrics.shuffleWriteMetrics =
       Utils.jsonOption(json \ "Shuffle Write Metrics").map(shuffleWriteMetricsFromJson)
+    metrics.inputMetrics =
+      Utils.jsonOption(json \ "Input Metrics").map(inputMetricsFromJson)
     metrics.updatedBlocks =
       Utils.jsonOption(json \ "Updated Blocks").map { value =>
         value.extract[List[JValue]].map { block =>
@@ -557,6 +568,13 @@ private[spark] object JsonProtocol {
     metrics
   }
 
+  def inputMetricsFromJson(json: JValue): InputMetrics = {
+    val metrics = new InputMetrics(
+      DataReadMethod.withName((json \ "Data Read Method").extract[String]))
+    metrics.bytesRead = (json \ "Bytes Read").extract[Long]
+    metrics
+  }
+
   def taskEndReasonFromJson(json: JValue): TaskEndReason = {
     val success = Utils.getFormattedClassName(Success)
     val resubmitted = Utils.getFormattedClassName(Resubmitted)

core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala

@@ -39,7 +39,11 @@ class JsonProtocolSuite extends FunSuite {
     val taskGettingResult =
       SparkListenerTaskGettingResult(makeTaskInfo(1000L, 2000, 5, 3000L, true))
     val taskEnd = SparkListenerTaskEnd(1, "ShuffleMapTask", Success,
-      makeTaskInfo(123L, 234, 67, 345L, false), makeTaskMetrics(300L, 400L, 500L, 600L, 700, 800))
+      makeTaskInfo(123L, 234, 67, 345L, false),
+      makeTaskMetrics(300L, 400L, 500L, 600L, 700, 800, hasHdfsInput = false))
+    val taskEndWithHdfsInput = SparkListenerTaskEnd(1, "ShuffleMapTask", Success,
+      makeTaskInfo(123L, 234, 345L),
+      makeTaskMetrics(300L, 400L, 500L, 600L, 700, 800, hasHdfsInput = true))
     val jobStart = SparkListenerJobStart(10, Seq[Int](1, 2, 3, 4), properties)
     val jobEnd = SparkListenerJobEnd(20, JobSucceeded)
     val environmentUpdate = SparkListenerEnvironmentUpdate(Map[String, Seq[(String, String)]](
@@ -61,6 +65,7 @@ class JsonProtocolSuite extends FunSuite {
     testEvent(taskStart, taskStartJsonString)
     testEvent(taskGettingResult, taskGettingResultJsonString)
     testEvent(taskEnd, taskEndJsonString)
+    testEvent(taskEndWithHdfsInput, taskEndWithHdfsInputJsonString)
     testEvent(jobStart, jobStartJsonString)
     testEvent(jobEnd, jobEndJsonString)
     testEvent(environmentUpdate, environmentUpdateJsonString)
@@ -74,8 +79,8 @@ class JsonProtocolSuite extends FunSuite {
   test("Dependent Classes") {
     testRDDInfo(makeRddInfo(2, 3, 4, 5L, 6L))
     testStageInfo(makeStageInfo(10, 20, 30, 40L, 50L))
-    testTaskInfo(makeTaskInfo(999L, 888, 55, 777L, false))
-    testTaskMetrics(makeTaskMetrics(33333L, 44444L, 55555L, 66666L, 7, 8))
+    testTaskInfo(makeTaskInfo(999L, 888, 777L, false))
+    testTaskMetrics(makeTaskMetrics(33333L, 44444L, 55555L, 66666L, 7, 8, hasHdfsInput = false))
     testBlockManagerId(BlockManagerId("Hong", "Kong", 500, 1000))
 
     // StorageLevel
@@ -118,7 +123,7 @@ class JsonProtocolSuite extends FunSuite {
     testBlockId(StreamBlockId(1, 2L))
   }
 
-  test("Backward compatibility") {
+  test("backward compatibility") {
     // StageInfo.details was added after 1.0.0.
     val info = makeStageInfo(1, 2, 3, 4L, 5L)
     assert(info.details.nonEmpty)
@@ -294,6 +299,8 @@ class JsonProtocolSuite extends FunSuite {
       metrics1.shuffleReadMetrics, metrics2.shuffleReadMetrics, assertShuffleReadEquals)
     assertOptionEquals(
       metrics1.shuffleWriteMetrics, metrics2.shuffleWriteMetrics, assertShuffleWriteEquals)
+    assertOptionEquals(
+      metrics1.inputMetrics, metrics2.inputMetrics, assertInputMetricsEquals)
     assertOptionEquals(metrics1.updatedBlocks, metrics2.updatedBlocks, assertBlocksEquals)
   }
 
@@ -311,6 +318,11 @@ class JsonProtocolSuite extends FunSuite {
     assert(metrics1.shuffleWriteTime === metrics2.shuffleWriteTime)
   }
 
+  private def assertEquals(metrics1: InputMetrics, metrics2: InputMetrics) {
+    assert(metrics1.readMethod === metrics2.readMethod)
+    assert(metrics1.bytesRead === metrics2.bytesRead)
+  }
+
   private def assertEquals(bm1: BlockManagerId, bm2: BlockManagerId) {
     assert(bm1.executorId === bm2.executorId)
     assert(bm1.host === bm2.host)
@@ -403,6 +415,10 @@ class JsonProtocolSuite extends FunSuite {
     assertEquals(w1, w2)
   }
 
+  private def assertInputMetricsEquals(i1: InputMetrics, i2: InputMetrics) {
+    assertEquals(i1, i2)
+  }
+
   private def assertTaskMetricsEquals(t1: TaskMetrics, t2: TaskMetrics) {
     assertEquals(t1, t2)
   }
@@ -460,9 +476,13 @@ class JsonProtocolSuite extends FunSuite {
     new TaskInfo(a, b, c, d, "executor", "your kind sir", TaskLocality.NODE_LOCAL, speculative)
   }
 
-  private def makeTaskMetrics(a: Long, b: Long, c: Long, d: Long, e: Int, f: Int) = {
+  /**
+   * Creates a TaskMetrics object describing a task that read data from HDFS (if hasHdfsInput is
+   * set to true) or read data from a shuffle otherwise.
+   */
+  private def makeTaskMetrics(
+      a: Long, b: Long, c: Long, d: Long, e: Int, f: Int, hasHdfsInput: Boolean) = {
     val t = new TaskMetrics
-    val sr = new ShuffleReadMetrics
     val sw = new ShuffleWriteMetrics
     t.hostname = "localhost"
     t.executorDeserializeTime = a
@@ -471,15 +491,23 @@ class JsonProtocolSuite extends FunSuite {
     t.jvmGCTime = d
     t.resultSerializationTime = a + b
     t.memoryBytesSpilled = a + c
-    sr.shuffleFinishTime = b + c
-    sr.totalBlocksFetched = e + f
-    sr.remoteBytesRead = b + d
-    sr.localBlocksFetched = e
-    sr.fetchWaitTime = a + d
-    sr.remoteBlocksFetched = f
+
+    if (hasHdfsInput) {
+      val inputMetrics = new InputMetrics(DataReadMethod.Hdfs)
+      inputMetrics.bytesRead = d + e + f
+      t.inputMetrics = Some(inputMetrics)
+    } else {
+      val sr = new ShuffleReadMetrics
+      sr.shuffleFinishTime = b + c
+      sr.totalBlocksFetched = e + f
+      sr.remoteBytesRead = b + d
+      sr.localBlocksFetched = e
+      sr.fetchWaitTime = a + d
+      sr.remoteBlocksFetched = f
+      t.shuffleReadMetrics = Some(sr)
+    }
     sw.shuffleBytesWritten = a + b + c
     sw.shuffleWriteTime = b + c + d
-    t.shuffleReadMetrics = Some(sr)
     t.shuffleWriteMetrics = Some(sw)
     // Make at most 6 blocks
     t.updatedBlocks = Some((1 to (e % 5 + 1)).map { i =>
@@ -568,6 +596,23 @@ class JsonProtocolSuite extends FunSuite {
       |}
     """.stripMargin
 
+  private val taskEndWithHdfsInputJsonString =
+    """
+      {"Event":"SparkListenerTaskEnd","Stage ID":1,"Task Type":"ShuffleMapTask",
+      "Task End Reason":{"Reason":"Success"},"Task Info":{"Task ID":123,"Index":
+      234,"Launch Time":345,"Executor ID":"executor","Host":"your kind sir",
+      "Locality":"NODE_LOCAL","Getting Result Time":0,"Finish Time":0,"Failed":
+      false,"Serialized Size":0},"Task Metrics":{"Host Name":"localhost",
+      "Executor Deserialize Time":300,"Executor Run Time":400,"Result Size":500,
+      "JVM GC Time":600,"Result Serialization Time":700,"Memory Bytes Spilled":
+      800,"Disk Bytes Spilled":0,"Shuffle Read Metrics":[],"Shuffle Write Metrics":
+      {"Shuffle Bytes Written":1200,"Shuffle Write Time":1500},"Input Metrics":
+      {"Data Read Method":Hdfs,"Bytes Read":1500},"Updated Blocks":
+      [{"Block ID":{"Type":"RDDBlockId","RDD ID":0,"Split Index":0},"Status":
+      {"Storage Level":{"Use Disk":true,"Use Memory":true,"Use Tachyon":false,
+      "Deserialized":false,"Replication":2},"Memory Size":0,"Disk Size":0,"Tachyon Size":0}}]}}
+    """
+
   private val jobStartJsonString =
     """
       {"Event":"SparkListenerJobStart","Job ID":10,"Stage IDs":[1,2,3,4],"Properties":