Fixed a lot of things

tdas · tdas · commit 22d7e6639125 · 2016-03-14T18:49:05.000-07:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -143,7 +143,7 @@ class HDFSBackedStateStoreProvider(
      */
     override def iterator(): Iterator[InternalRow] = {
       verify(state == COMMITTED, "Cannot get iterator of store data before comitting")
-      HDFSBackedStateStoreProvider.this.iterator(version)
+      HDFSBackedStateStoreProvider.this.iterator(newVersion)
     }
 
     /**
@@ -245,10 +245,8 @@ class HDFSBackedStateStoreProvider(
 
   /** Load the required version of the map data from the backing files */
   private def loadMap(version: Long): MapType = {
-    if (version < 0) return new MapType
-    synchronized {
-      loadedMaps.get(version)
-    }.getOrElse {
+    if (version <= 0) return new MapType
+    synchronized { loadedMaps.get(version) }.getOrElse {
       val mapFromFile = readSnapshotFile(version).getOrElse {
         val prevMap = loadMap(version - 1)
         val deltaUpdates = readDeltaFile(version)
@@ -328,9 +326,7 @@ class HDFSBackedStateStoreProvider(
         val lastVersion = files.last.version
         val deltaFilesForLastVersion =
           filesForVersion(files, lastVersion).filter(_.isSnapshot == false)
-        synchronized {
-          loadedMaps.get(lastVersion)
-        } match {
+        synchronized { loadedMaps.get(lastVersion) } match {
           case Some(map) =>
             if (deltaFilesForLastVersion.size > maxDeltaChainForSnapshots) {
               writeSnapshotFile(lastVersion, map)
@@ -342,7 +338,7 @@ class HDFSBackedStateStoreProvider(
       }
     } catch {
       case NonFatal(e) =>
-        logWarning(s"Error doing snapshots for $this")
+        logWarning(s"Error doing snapshots for $this", e)
     }
   }
 
@@ -356,7 +352,7 @@ class HDFSBackedStateStoreProvider(
       val files = fetchFiles()
       if (files.nonEmpty) {
         val earliestVersionToRetain = files.last.version - numBatchesToRetain
-        if (earliestVersionToRetain >= 0) {
+        if (earliestVersionToRetain > 0) {
           val earliestFileToRetain = filesForVersion(files, earliestVersionToRetain).head
           synchronized {
             loadedMaps.keys.filter(_ < earliestVersionToRetain).foreach(loadedMaps.remove)
@@ -368,7 +364,7 @@ class HDFSBackedStateStoreProvider(
       }
     } catch {
       case NonFatal(e) =>
-        logWarning(s"Error cleaning up files for $this")
+        logWarning(s"Error cleaning up files for $this", e)
     }
   }
 
@@ -380,7 +376,6 @@ class HDFSBackedStateStoreProvider(
       .filter(_.isSnapshot == true)
       .takeWhile(_.version <= version)
       .lastOption
-
     val deltaBatchFiles = latestSnapshotFileBeforeVersion match {
       case Some(snapshotFile) =>
         val deltaBatchIds = (snapshotFile.version + 1) to version
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreRDD.scala
@@ -30,10 +30,12 @@ class StateStoreRDD[INPUT: ClassTag, OUTPUT: ClassTag](
     dataRDD: RDD[INPUT],
     storeUpdateFunction: (StateStore, Iterator[INPUT]) => Iterator[OUTPUT],
     operatorId: Long,
-    newStoreVersion: Long,
+    storeVersion: Long,
     storeDirectory: String,
     storeCoordinator: StateStoreCoordinator) extends RDD[OUTPUT](dataRDD) {
 
+  val nextVersion = storeVersion + 1
+
   override protected def getPartitions: Array[Partition] = dataRDD.partitions
   override def getPreferredLocations(partition: Partition): Seq[String] = {
     Seq.empty
@@ -47,11 +49,8 @@ class StateStoreRDD[INPUT: ClassTag, OUTPUT: ClassTag](
     var store: StateStore = null
 
     Utils.tryWithSafeFinally {
-      store = StateStore.get(
-        StateStoreId(operatorId, partition.index),
-        storeDirectory,
-        newStoreVersion - 1
-      )
+      val storeId = StateStoreId(operatorId, partition.index)
+      store = StateStore.get(storeId, storeDirectory, storeVersion)
       val inputIter = dataRDD.compute(partition, ctxt)
       val outputIter = storeUpdateFunction(store, inputIter)
       assert(store.hasCommitted)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
@@ -27,12 +27,12 @@ package object state {
     def withStateStores[OUTPUT: ClassTag](
       storeUpdateFunction: (StateStore, Iterator[INPUT]) => Iterator[OUTPUT],
       operatorId: Long,
-      newStoreVersion: Long,
+      storeVersion: Long,
       storeDirectory: String,
       storeCoordinator: StateStoreCoordinator
-    ): RDD[OUTPUT] = {
+    ): StateStoreRDD[INPUT, OUTPUT] = {
       new StateStoreRDD(
-        dataRDD, storeUpdateFunction, operatorId, newStoreVersion, storeDirectory, storeCoordinator)
+        dataRDD, storeUpdateFunction, operatorId, storeVersion, storeDirectory, storeCoordinator)
     }
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreRDDSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreRDDSuite.scala
@@ -17,6 +17,9 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
+import java.io.File
+import java.nio.file.Files
+
 import scala.util.Random
 
 import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll}
@@ -29,15 +32,21 @@ import org.apache.spark.util.Utils
 class StateStoreRDDSuite extends SparkFunSuite with BeforeAndAfter with BeforeAndAfterAll {
 
   private val conf = new SparkConf().setMaster("local").setAppName(this.getClass.getCanonicalName)
-  private var tempDir = Utils.createTempDir().toString
+  private var tempDir = Files.createTempDirectory("StateStoreRDDSuite").toString
+  println(tempDir)
 
   import StateStoreSuite._
 
   after {
     StateStore.clearAll()
   }
 
-  test("versioning and immuability") {
+  override def afterAll(): Unit = {
+    super.afterAll()
+    Utils.deleteRecursively(new File(tempDir))
+  }
+
+  test("versioning and immutability") {
     withSpark(new SparkContext(conf)) { sc =>
       val path = Utils.createDirectory(tempDir, Random.nextString(10)).toString
       val increment = (store: StateStore, iter: Iterator[String]) => {
@@ -53,20 +62,53 @@ class StateStoreRDDSuite extends SparkFunSuite with BeforeAndAfter with BeforeAn
       }
       val opId = 0
       val rdd1 = makeRDD(sc, Seq("a", "b", "a"))
-        .withStateStores(increment, opId, newStoreVersion = 0, path, null)
+        .withStateStores(increment, opId, storeVersion = 0, path, null)
       assert(rdd1.collect().toSet === Set("a" -> 2, "b" -> 1))
 
       // Generate next version of stores
       val rdd2 = makeRDD(sc, Seq("a", "c"))
-        .withStateStores(increment, opId, newStoreVersion = 1, path, null)
+        .withStateStores(increment, opId, storeVersion = 1, path, null)
       assert(rdd2.collect().toSet === Set("a" -> 3, "b" -> 1, "c" -> 1))
 
       // Make sure the previous RDD still has the same data.
       assert(rdd1.collect().toSet === Set("a" -> 2, "b" -> 1))
     }
   }
 
+  test("recovering from files") {
+    val opId = 0
+    val path = Utils.createDirectory(tempDir, Random.nextString(10)).toString
+
+    def makeStoreRDD(sc: SparkContext, seq: Seq[String], storeVersion: Int): RDD[(String, Int)] = {
+      makeRDD(sc, Seq("a")).withStateStores(increment, opId, storeVersion, path, null)
+    }
+
+    // Generate RDDs and state store data
+    withSpark(new SparkContext(conf)) { sc =>
+      for (i <- 1 to 20) {
+        require(makeStoreRDD(sc, Seq("a"), i - 1).collect().toSet === Set("a" -> i))
+      }
+    }
+
+    // With a new context, try using the earlier state store data
+    withSpark(new SparkContext(conf)) { sc =>
+      assert(makeStoreRDD(sc, Seq("a"), 20).collect().toSet === Set("a" -> 21))
+    }
+  }
+
   private def makeRDD(sc: SparkContext, seq: Seq[String]): RDD[String] = {
     sc.makeRDD(seq, 2).groupBy(x => x).flatMap(_._2)
   }
+
+  private val increment = (store: StateStore, iter: Iterator[String]) => {
+    iter.foreach { s =>
+      store.update(
+        wrapKey(s), oldRow => {
+          val oldValue = oldRow.map(unwrapValue).getOrElse(0)
+          wrapValue(oldValue + 1)
+        })
+    }
+    store.commit()
+    store.iterator().map(unwrapKeyValue)
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
@@ -74,9 +74,8 @@ class StateStoreSuite extends SparkFunSuite with BeforeAndAfter with PrivateMeth
     assert(store.commit() === 1)
 
     assert(store.hasCommitted)
-    assert(store.iterator() === Set("b" -> 2))
-    assert(store.updates() === Set("b" -> 2))
-    assert(provider.latestIterator() === Set("b" -> 2))
+    assert(unwrapToSet(store.iterator()) === Set("b" -> 2))
+    assert(unwrapToSet(provider.latestIterator()) === Set("b" -> 2))
     assert(fileExists(provider, version = 1, isSnapshot = false))
     assert(getDataFromFiles(provider) === Set("b" -> 2))
 
@@ -92,8 +91,7 @@ class StateStoreSuite extends SparkFunSuite with BeforeAndAfter with PrivateMeth
     val reloadedStore = new HDFSBackedStateStoreProvider(store.id, provider.directory).getStore(1)
     update(reloadedStore, "c", 4)
     assert(reloadedStore.commit() === 2)
-    assert(reloadedStore.iterator() === Set("b" -> 2, "c" -> 4))
-    assert(reloadedStore.updates() === Set("c" -> 4))
+    assert(unwrapToSet(reloadedStore.iterator()) === Set("b" -> 2, "c" -> 4))
     assert(getDataFromFiles(provider) === Set("b" -> 2, "c" -> 4))
     assert(getDataFromFiles(provider, version = 1) === Set("b" -> 2))
     assert(getDataFromFiles(provider, version = 2) === Set("b" -> 2, "c" -> 4))
@@ -104,7 +102,7 @@ class StateStoreSuite extends SparkFunSuite with BeforeAndAfter with PrivateMeth
     val store = provider.getStore(0)
     update(store, "a", 1)
     store.commit()
-    assert(store.iterator() === Set("a" -> 1))
+    assert(unwrapToSet(store.iterator()) === Set("a" -> 1))
 
     // cancelUpdates should not change the data in the files
     val store1 = provider.getStore(1)
@@ -124,87 +122,92 @@ class StateStoreSuite extends SparkFunSuite with BeforeAndAfter with PrivateMeth
     val store = provider.getStore(0)
     update(store, "a", 1)
     assert(store.commit() === 1)
-    assert(store.iterator() === Set("a" -> 1))
+    assert(unwrapToSet(store.iterator()) === Set("a" -> 1))
 
     intercept[IllegalStateException] {
       provider.getStore(2)
     }
 
     // Update store version with some data
-    provider.getStore(1)
-    update(store, "b", 1)
-    assert(store.commit() === 2)
-    assert(store.iterator() === Set("a" -> 1, "b" -> 1))
+    val store1 = provider.getStore(1)
+    update(store1, "b", 1)
+    assert(store1.commit() === 2)
+    assert(unwrapToSet(store1.iterator()) === Set("a" -> 1, "b" -> 1))
     assert(getDataFromFiles(provider) === Set("a" -> 1, "b" -> 1))
 
     // Overwrite the version with other data
-    provider.getStore(1)
-    update(store, "c", 1)
-    assert(store.commit() === 2)
-    assert(store.iterator() === Set("a" -> 1, "c" -> 1))
+    val store2 = provider.getStore(1)
+    update(store2, "c", 1)
+    assert(store2.commit() === 2)
+    assert(unwrapToSet(store2.iterator()) === Set("a" -> 1, "c" -> 1))
     assert(getDataFromFiles(provider) === Set("a" -> 1, "c" -> 1))
   }
 
   test("snapshotting") {
     val provider = newStoreProvider(maxDeltaChainForSnapshots = 5)
 
-    var currentVersion = -1
+    var currentVersion = 0
     def updateVersionTo(targetVersion: Int): Unit = {
       for (i <- currentVersion + 1 to targetVersion) {
-        val store = provider.getStore(i - 1)
+        val store = provider.getStore(currentVersion)
         update(store, "a", i)
         store.commit()
-
+        currentVersion += 1
       }
-      currentVersion = targetVersion
+      require(currentVersion === targetVersion)
     }
 
+
     updateVersionTo(2)
     require(getDataFromFiles(provider) === Set("a" -> 2))
     provider.manage()               // should not generate snapshot files
     assert(getDataFromFiles(provider) === Set("a" -> 2))
-    for (i <- 0 to 2) {
+
+    for (i <- 1 to currentVersion) {
       assert(fileExists(provider, i, isSnapshot = false))  // all delta files present
       assert(!fileExists(provider, i, isSnapshot = true))  // no snapshot files present
     }
 
     // After version 6, snapshotting should generate one snapshot file
     updateVersionTo(6)
-    require(getDataFromFiles(provider) === Set("a" -> 6), "Store not updated correctly")
+    require(getDataFromFiles(provider) === Set("a" -> 6), "store not updated correctly")
     provider.manage()       // should generate snapshot files
     assert(getDataFromFiles(provider) === Set("a" -> 6), "snapshotting messed up the data")
     assert(getDataFromFiles(provider) === Set("a" -> 6))
 
     val snapshotVersion = (0 to 6).find(version => fileExists(provider, version, isSnapshot = true))
-    assert(snapshotVersion.nonEmpty, "Snapshot file not generated")
-
+    assert(snapshotVersion.nonEmpty, "snapshot file not generated")
 
     // After version 20, snapshotting should generate newer snapshot files
     updateVersionTo(20)
-    require(getDataFromFiles(provider) === Set("a" -> 20), "Store not updated correctly")
+    require(getDataFromFiles(provider) === Set("a" -> 20), "store not updated correctly")
     provider.manage()       // do snapshot
     assert(getDataFromFiles(provider) === Set("a" -> 20), "snapshotting messed up the data")
     assert(getDataFromFiles(provider) === Set("a" -> 20))
 
     val latestSnapshotVersion = (0 to 20).filter(version =>
       fileExists(provider, version, isSnapshot = true)).lastOption
-    assert(latestSnapshotVersion.nonEmpty, "No snapshot file found")
-    assert(latestSnapshotVersion.get > snapshotVersion.get, "Newer snapshot not generated")
+    assert(latestSnapshotVersion.nonEmpty, "no snapshot file found")
+    assert(latestSnapshotVersion.get > snapshotVersion.get, "newer snapshot not generated")
 
   }
 
   test("cleaning") {
     val provider = newStoreProvider(maxDeltaChainForSnapshots = 5)
 
-    for (i <- 0 to 20) {
-      val store = provider.getStore(i)
+    for (i <- 1 to 20) {
+      val store = provider.getStore(i - 1)
       update(store, "a", i)
       store.commit()
+      provider.manage() // do cleanup
     }
-    require(provider.latestIterator() === Set("a" -> 20), "Store not updated correctly")
-    provider.manage()     // do cleanup
-    assert(fileExists(provider, 0, isSnapshot = false))
+    require(
+      unwrapToSet(provider.latestIterator()) === Set("a" -> 20),
+      "store not updated correctly")
+
+    assert(!fileExists(provider, version = 1, isSnapshot = false)) // first file should be deleted
 
+    // last couple of versions should be retrievable
     assert(getDataFromFiles(provider, 20) === Set("a" -> 20))
     assert(getDataFromFiles(provider, 19) === Set("a" -> 19))
   }
@@ -214,7 +217,7 @@ class StateStoreSuite extends SparkFunSuite with BeforeAndAfter with PrivateMeth
     version: Int = -1): Set[(String, Int)] = {
     val reloadedProvider = new HDFSBackedStateStoreProvider(provider.id, provider.directory)
     if (version < 0) {
-      reloadedProvider.latestIterator.map(unwrapKeyValue).toSet
+      reloadedProvider.latestIterator().map(unwrapKeyValue).toSet
     } else {
       reloadedProvider.iterator(version).map(unwrapKeyValue).toSet
     }
@@ -295,4 +298,8 @@ private[state] object StateStoreSuite {
   def unwrapKeyValue(row: InternalRow): (String, Int) = {
     (row.getString(0), row.getInt(1))
   }
+
+  def unwrapToSet(iterator: Iterator[InternalRow]): Set[(String, Int)] = {
+    iterator.map(unwrapKeyValue).toSet
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -27,12 +27,12 @@ package object state {`
`27`	`27`	`def withStateStores[OUTPUT: ClassTag](`
`28`	`28`	`storeUpdateFunction: (StateStore, Iterator[INPUT]) => Iterator[OUTPUT],`
`29`	`29`	`operatorId: Long,`
`30`		`- newStoreVersion: Long,`
	`30`	`+ storeVersion: Long,`
`31`	`31`	`storeDirectory: String,`
`32`	`32`	`storeCoordinator: StateStoreCoordinator`
`33`		`- ): RDD[OUTPUT] = {`
	`33`	`+ ): StateStoreRDD[INPUT, OUTPUT] = {`
`34`	`34`	`new StateStoreRDD(`
`35`		`- dataRDD, storeUpdateFunction, operatorId, newStoreVersion, storeDirectory, storeCoordinator)`
	`35`	`+ dataRDD, storeUpdateFunction, operatorId, storeVersion, storeDirectory, storeCoordinator)`
`36`	`36`	`}`
`37`	`37`	`}`
`38`	`38`	`}`