Fix JsonProtocol compatibility

Author: JoshRosen

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

@@ -31,6 +31,21 @@ import org.apache.spark.scheduler._
 import org.apache.spark.storage._
 import org.apache.spark._
 
+/**
+ * Serializes SparkListener events to/from JSON.  This protocol provides strong backwards-
+ * and forwards-compatibility guarantees: any version of Spark should be able to read JSON output
+ * written by any other version, including newer versions.
+ *
+ * JsonProtocolSuite contains backwards-compatibility tests which check that the current version of
+ * JsonProtocol is able to read output written by earlier versions.  We do not currently have tests
+ * for reading newer JSON output with older Spark versions.
+ *
+ * To ensure that we provide these guarantees, follow these rules when modifying these methods:
+ *
+ *  - Never delete any JSON fields.
+ *  - Any new JSON fields should be optional; use `Utils.jsonOption` when reading these fields
+ *    in `*FromJson` methods.
+ */
 private[spark] object JsonProtocol {
   // TODO: Remove this file and put JSON serialization into each individual class.
 
@@ -121,8 +136,8 @@ private[spark] object JsonProtocol {
     val properties = propertiesToJson(jobStart.properties)
     ("Event" -> Utils.getFormattedClassName(jobStart)) ~
     ("Job ID" -> jobStart.jobId) ~
-    // ("Stage IDs" -> jobStart.stageIds) ~  // Removed in 1.2.0
-    ("Stage Infos" -> jobStart.stageInfos.map(stageInfoToJson)) ~  // Added in 1.2.0
+    ("Stage Infos" -> jobStart.stageInfos.map(stageInfoToJson)) ~  // Added in Spark 1.2.0
+    ("Stage IDs" -> jobStart.stageIds) ~
     ("Properties" -> properties)
   }
 
@@ -454,20 +469,13 @@ private[spark] object JsonProtocol {
 
   def jobStartFromJson(json: JValue): SparkListenerJobStart = {
     val jobId = (json \ "Job ID").extract[Int]
+    val stageIds = (json \ "Stage IDs").extract[List[JValue]].map(_.extract[Int])
     val properties = propertiesFromJson(json \ "Properties")
-    val stageInfos = {
-      // Prior to 1.2.0, we serialized stageIds but not stageInfos; in 1.2.0, we do the opposite.
-      // This block of code handles backwards compatibility:
-      val stageIds: Option[Seq[Int]] =
-        Utils.jsonOption(json \ "Stage IDs").map(_.extract[List[JValue]].map(_.extract[Int]))
-      if (stageIds.isDefined) { // Reading JSON written prior to 1.2.0
-        stageIds.get.map(id => new StageInfo(id, 0, "unknown", 0, Seq.empty, "unknown"))
-      } else { // Reading JSON written after 1.2.0
-        Utils.jsonOption(json \ "Stage Infos")
-          .map(_.extract[Seq[JValue]].map(stageInfoFromJson)).getOrElse(Seq.empty)
+    // The "Stage Infos" field was added in Spark 1.2.0
+    val stageInfos = Utils.jsonOption(json \ "Stage Infos")
+      .map(_.extract[Seq[JValue]].map(stageInfoFromJson)).getOrElse {
+        stageIds.map(id => new StageInfo(id, 0, "unknown", 0, Seq.empty, "unknown"))
       }
-    }
-
     SparkListenerJobStart(jobId, stageInfos, properties)
   }
 

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

@@ -21,8 +21,6 @@ import java.util.Properties
 
 import scala.collection.Map
 
-import org.json4s.JsonAST.JObject
-import org.json4s.JsonDSL._
 import org.json4s.jackson.JsonMethods._
 import org.scalatest.FunSuite
 
@@ -51,7 +49,8 @@ class JsonProtocolSuite extends FunSuite {
       makeTaskMetrics(300L, 400L, 500L, 600L, 700, 800, hasHadoopInput = true, hasOutput = true))
     val jobStart = {
       val stageIds = Seq[Int](1, 2, 3, 4)
-      val stageInfos = stageIds.map(x => makeStageInfo(x, x * 200, x * 300, x * 400L, x * 500L))
+      val stageInfos = stageIds.map(x =>
+        makeStageInfo(x, x * 200, x * 300, x * 400L, x * 500L))
       SparkListenerJobStart(10, stageInfos, properties)
     }
     val jobEnd = SparkListenerJobEnd(20, JobSucceeded)
@@ -233,11 +232,13 @@ class JsonProtocolSuite extends FunSuite {
   test("SparkListenerJobStart backward compatibility") {
     // Prior to Spark 1.2.0, SparkListenerJobStart did not have a "Stage Infos" property.
     val stageIds = Seq[Int](1, 2, 3, 4)
-    val stageInfos = stageIds.map(id => new StageInfo(id, 0, "unknown", 0, Seq.empty, "unknown"))
+    val stageInfos = stageIds.map(x => makeStageInfo(x, x * 200, x * 300, x * 400, x * 500))
+    val dummyStageInfos =
+      stageIds.map(id => new StageInfo(id, 0, "unknown", 0, Seq.empty, "unknown"))
     val jobStart = SparkListenerJobStart(10, stageInfos, properties)
     val oldEvent = JsonProtocol.jobStartToJson(jobStart).removeField({_._1 == "Stage Infos"})
-      .asInstanceOf[JObject] ~ ("Stage IDs" -> stageIds)
-    val expectedJobStart = SparkListenerJobStart(10, stageInfos, properties)
+    val expectedJobStart =
+      SparkListenerJobStart(10, dummyStageInfos, properties)
     assertEquals(expectedJobStart, JsonProtocol.jobStartFromJson(oldEvent))
   }
 
@@ -1322,6 +1323,12 @@ class JsonProtocolSuite extends FunSuite {
       |      ]
       |    }
       |  ],
+      |  "Stage IDs": [
+      |    1,
+      |    2,
+      |    3,
+      |    4
+      |  ],
       |  "Properties": {
       |    "France": "Paris",
       |    "Germany": "Berlin",