Update tests. Remove var in OffsetSeqMetadata.

Author: kunalkhamar

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala

@@ -72,8 +72,8 @@ object OffsetSeq {
  * @param conf: Additional conf_s to be persisted across batches, e.g. number of shuffle partitions.
  */
 case class OffsetSeqMetadata(
-    var batchWatermarkMs: Long = 0,
-    var batchTimestampMs: Long = 0,
+    batchWatermarkMs: Long = 0,
+    batchTimestampMs: Long = 0,
     conf: Map[String, String] = Map.empty) {
   def json: String = Serialization.write(this)(OffsetSeqMetadata.format)
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala

@@ -464,25 +464,28 @@ class StreamExecution(
       }
     }
     if (hasNewData) {
-      // Current batch timestamp in milliseconds
-      offsetSeqMetadata.batchTimestampMs = triggerClock.getTimeMillis()
+      var batchWatermarkMs = offsetSeqMetadata.batchWatermarkMs
       // Update the eventTime watermark if we find one in the plan.
       if (lastExecution != null) {
         lastExecution.executedPlan.collect {
           case e: EventTimeWatermarkExec if e.eventTimeStats.value.count > 0 =>
             logDebug(s"Observed event time stats: ${e.eventTimeStats.value}")
             e.eventTimeStats.value.max - e.delayMs
         }.headOption.foreach { newWatermarkMs =>
-          if (newWatermarkMs > offsetSeqMetadata.batchWatermarkMs) {
+          if (newWatermarkMs > batchWatermarkMs) {
             logInfo(s"Updating eventTime watermark to: $newWatermarkMs ms")
-            offsetSeqMetadata.batchWatermarkMs = newWatermarkMs
+            batchWatermarkMs = newWatermarkMs
           } else {
             logDebug(
               s"Event time didn't move: $newWatermarkMs < " +
-                s"${offsetSeqMetadata.batchWatermarkMs}")
+                s"$batchWatermarkMs")
           }
         }
       }
+      offsetSeqMetadata = OffsetSeqMetadata(
+        batchWatermarkMs,
+        triggerClock.getTimeMillis(), // Current batch timestamp in milliseconds
+        offsetSeqMetadata.conf) // Keep the same conf
 
       updateStatusMessage("Writing offsets to log")
       reportTimeTaken("walCommit") {

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala

@@ -23,6 +23,8 @@ import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit}
 import scala.reflect.ClassTag
 import scala.util.control.ControlThrowable
 
+import org.apache.commons.io.FileUtils
+
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.streaming.InternalOutputModes
 import org.apache.spark.sql.execution.command.ExplainCommand
@@ -409,41 +411,88 @@ class StreamSuite extends StreamTest {
       CheckAnswer((1, 2), (2, 2), (3, 2)))
   }
 
-  test("SPARK-19873: backward compat with checkpoints that do not record shuffle partitions") {
-    val inputData = MemoryStream[Int]
-    inputData.addData(1, 2, 3, 4)
-    inputData.addData(3, 4, 5, 6)
-    inputData.addData(5, 6, 7, 8)
+  test("SPARK-19873: backward compatibility - recover from a Spark v2.1 checkpoint") {
+    var inputData: MemoryStream[Int] = null
+    var query: DataStreamWriter[Row] = null
 
-    val resourceUri =
-      this.getClass.getResource("/structured-streaming/checkpoint-version-2.1.0").toURI
-    val checkpointDir = new File(resourceUri).getCanonicalPath
-    val query = inputData
-      .toDF()
-      .groupBy($"value")
-      .agg(count("*"))
-      .writeStream
-      .queryName("counts")
-      .outputMode("complete")
-      .option("checkpointLocation", checkpointDir)
-      .format("memory")
-
-    // Checkpoint data was generated by a query with 10 shuffle partitions.
-    // Test if recovery from checkpoint is successful.
-    withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
-      query.start().processAllAvailable()
-
-      QueryTest.checkAnswer(spark.table("counts").toDF(),
-        Row("1", 1) :: Row("2", 1) :: Row("3", 2) :: Row("4", 2) ::
-        Row("5", 2) :: Row("6", 2) :: Row("7", 1) :: Row("8", 1) :: Nil)
+    def init(): Unit = {
+      inputData = MemoryStream[Int]
+      inputData.addData(1, 2, 3, 4)
+      inputData.addData(3, 4, 5, 6)
+      inputData.addData(5, 6, 7, 8)
+
+      query = inputData
+        .toDF()
+        .groupBy($"value")
+        .agg(count("*"))
+        .writeStream
+        .outputMode("complete")
+        .format("memory")
     }
 
-    // If the number of partitions is greater, should throw exception.
-    withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "15") {
-      intercept[IllegalArgumentException] {
-        query.start().processAllAvailable()
+    // Get an existing checkpoint generated by Spark v2.1.
+    // v2.1 does not record # shuffle partitions in the offset metadata.
+    val resourceUri =
+      this.getClass.getResource("/structured-streaming/checkpoint-version-2.1.0").toURI
+    val checkpointDir = new File(resourceUri)
+
+    // 1 - Test if recovery from the checkpoint is successful.
+    init()
+    withTempDir(dir => {
+      // Copy the checkpoint to a temp dir to prevent changes to the original.
+      // Not doing this will lead to the test passing on the first run, but fail subsequent runs.
+      FileUtils.copyDirectory(checkpointDir, dir)
+
+      // Checkpoint data was generated by a query with 10 shuffle partitions.
+      // In order to test reading from the checkpoint, the checkpoint must have two or more batches,
+      // since the last batch may be rerun.
+      withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
+        var streamingQuery: StreamingQuery = null
+        try {
+          streamingQuery =
+            query
+              .queryName("counts")
+              .option("checkpointLocation", dir.getCanonicalPath)
+              .start()
+          streamingQuery.processAllAvailable()
+          inputData.addData(9)
+          streamingQuery.processAllAvailable()
+
+          QueryTest.checkAnswer(spark.table("counts").toDF(),
+            Row("1", 1) :: Row("2", 1) :: Row("3", 2) :: Row("4", 2) ::
+            Row("5", 2) :: Row("6", 2) :: Row("7", 1) :: Row("8", 1) :: Row("9", 1) :: Nil)
+        } finally {
+          if (streamingQuery ne null) {
+            streamingQuery.stop()
+          }
+        }
       }
-    }
+    })
+
+    // 2 - Check recovery with wrong num shuffle partitions
+    init()
+    withTempDir(dir => {
+      FileUtils.copyDirectory(checkpointDir, dir)
+
+      // Since the number of partitions is greater than 10, should throw exception.
+      withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "15") {
+        var streamingQuery: StreamingQuery = null
+        try {
+          intercept[StreamingQueryException] {
+            streamingQuery =
+              query
+                .queryName("badQuery")
+                .option("checkpointLocation", dir.getCanonicalPath)
+                .start()
+            streamingQuery.processAllAvailable()
+          }
+        } finally {
+          if (streamingQuery ne null) {
+            streamingQuery.stop()
+          }
+        }
+      }
+    })
   }
 }
 

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala

@@ -152,12 +152,10 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
       AssertOnQuery(q => {
         q.exception.get.startOffset ===
           q.committedOffsets.toOffsetSeq(
-            Seq(inputData),
-            OffsetSeqMetadata(0, 0)).toString &&
+            Seq(inputData), OffsetSeqMetadata()).toString &&
           q.exception.get.endOffset ===
             q.availableOffsets.toOffsetSeq(
-              Seq(inputData),
-              OffsetSeqMetadata(0, 0)).toString
+              Seq(inputData), OffsetSeqMetadata()).toString
       }, "incorrect start offset or end offset on exception")
     )
   }