[SPARK-13809][SQL] State store for streaming aggregations

sql/core/src/main/scala/org/apache/spark/sql/ContinuousQueryManager.scala

@@ -34,7 +34,8 @@ import org.apache.spark.sql.util.ContinuousQueryListener
 @Experimental
 class ContinuousQueryManager(sqlContext: SQLContext) {
 
-  private[sql] val stateStoreCoordinator = StateStoreCoordinatorRef(sqlContext.sparkContext.env)
+  private[sql] val stateStoreCoordinator =
+    StateStoreCoordinatorRef.forDriver(sqlContext.sparkContext.env)
   private val listenerBus = new ContinuousQueryListenerBus(sqlContext.sparkContext.listenerBus)
   private val activeQueries = new mutable.HashMap[String, ContinuousQuery]
   private val activeQueriesLock = new Object

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

@@ -583,6 +583,3 @@ private[state] class HDFSBackedStateStoreProvider(
   }
 }
 
-private[state] object HDFSBackedStateStoreProvider {
-
-}

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

@@ -35,6 +35,7 @@ import org.apache.spark.util.ThreadUtils
 /** Unique identifier for a [[StateStore]] */
 case class StateStoreId(checkpointLocation: String, operatorId: Long, partitionId: Int)
 
+
 /**
  * Base trait for a versioned key-value store used for streaming aggregations
  */
@@ -84,6 +85,7 @@ trait StateStore {
   def hasCommitted: Boolean
 }
 
+
 /** Trait representing a provider of a specific version of a [[StateStore]]. */
 trait StateStoreProvider {
 
@@ -94,6 +96,7 @@ trait StateStoreProvider {
   def doMaintenance(): Unit = { }
 }
 
+
 /** Trait representing updates made to a [[StateStore]]. */
 sealed trait StoreUpdate
 
@@ -118,7 +121,6 @@ private[state] object StateStore extends Logging {
   val MAINTENANCE_INTERVAL_DEFAULT_SECS = 60
 
   private val loadedProviders = new mutable.HashMap[StateStoreId, StateStoreProvider]()
-  private val maintenanceTimer = new Timer("StateStore Timer", true)
   private val maintenanceTaskExecutor =
     ThreadUtils.newDaemonSingleThreadScheduledExecutor("state-store-maintenance-task")
 
@@ -232,7 +234,7 @@ private[state] object StateStore extends Logging {
     val env = SparkEnv.get
     if (env != null) {
       if (_coordRef == null) {
-        _coordRef = StateStoreCoordinatorRef(env)
+        _coordRef = StateStoreCoordinatorRef.forExecutor(env)
       }
       logDebug(s"Retrieved reference to StateStoreCoordinator: ${_coordRef}")
       Some(_coordRef)

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

@@ -53,19 +53,25 @@ private[sql] object StateStoreCoordinatorRef extends Logging {
    * Create a reference to a [[StateStoreCoordinator]], This can be called from driver as well as
    * executors.
    */
-  def apply(env: SparkEnv): StateStoreCoordinatorRef = synchronized {
+  def forDriver(env: SparkEnv): StateStoreCoordinatorRef = synchronized {
     try {
       val coordinator = new StateStoreCoordinator(env.rpcEnv)
       val coordinatorRef = env.rpcEnv.setupEndpoint(endpointName, coordinator)
       logInfo("Registered StateStoreCoordinator endpoint")
       new StateStoreCoordinatorRef(coordinatorRef)
     } catch {
       case e: IllegalArgumentException =>
-        logDebug("Retrieving existing StateStoreCoordinator endpoint")
         val rpcEndpointRef = RpcUtils.makeDriverRef(endpointName, env.conf, env.rpcEnv)
+        logDebug("Retrieved existing StateStoreCoordinator endpoint")
         new StateStoreCoordinatorRef(rpcEndpointRef)
     }
   }
+
+  def forExecutor(env: SparkEnv): StateStoreCoordinatorRef = synchronized {
+    val rpcEndpointRef = RpcUtils.makeDriverRef(endpointName, env.conf, env.rpcEnv)
+    logDebug("Retrieved existing StateStoreCoordinator endpoint")
+    new StateStoreCoordinatorRef(rpcEndpointRef)
+  }
 }
 
 /**

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreCoordinatorSuite.scala

@@ -94,7 +94,7 @@ class StateStoreCoordinatorSuite extends SparkFunSuite with SharedSparkContext {
 
   test("multiple references have same underlying coordinator") {
     withCoordinatorRef(sc) { coordRef1 =>
-      val coordRef2 = StateStoreCoordinatorRef(sc.env)
+      val coordRef2 = StateStoreCoordinatorRef.forDriver(sc.env)
 
       val id = StateStoreId("x", 0, 0)
 
@@ -114,7 +114,7 @@ object StateStoreCoordinatorSuite {
   def withCoordinatorRef(sc: SparkContext)(body: StateStoreCoordinatorRef => Unit): Unit = {
     var coordinatorRef: StateStoreCoordinatorRef = null
     try {
-      coordinatorRef = StateStoreCoordinatorRef(sc.env)
+      coordinatorRef = StateStoreCoordinatorRef.forDriver(sc.env)
       body(coordinatorRef)
     } finally {
       if (coordinatorRef != null) coordinatorRef.stop()

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreRDDSuite.scala

@@ -142,6 +142,38 @@ class StateStoreRDDSuite extends SparkFunSuite with BeforeAndAfter with BeforeAn
     }
   }
 
+  test("distributed test") {
+    quietly {
+      withSpark(new SparkContext(sparkConf.setMaster("local-cluster[2, 1, 1024]"))) { sc =>
+        implicit val sqlContet = new SQLContext(sc)
+        val path = Utils.createDirectory(tempDir, Random.nextString(10)).toString
+        val increment = (store: StateStore, iter: Iterator[String]) => {
+          iter.foreach { s =>
+            store.update(
+              stringToRow(s), oldRow => {
+                val oldValue = oldRow.map(rowToInt).getOrElse(0)
+                intToRow(oldValue + 1)
+              })
+          }
+          store.commit()
+          store.iterator().map(rowsToStringInt)
+        }
+        val opId = 0
+        val rdd1 = makeRDD(sc, Seq("a", "b", "a")).mapPartitionWithStateStore(
+          increment, path, opId, storeVersion = 0, keySchema, valueSchema)
+        assert(rdd1.collect().toSet === Set("a" -> 2, "b" -> 1))
+
+        // Generate next version of stores
+        val rdd2 = makeRDD(sc, Seq("a", "c")).mapPartitionWithStateStore(
+          increment, path, opId, storeVersion = 1, keySchema, valueSchema)
+        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))
+      }
+    }
+  }
+
   private def makeRDD(sc: SparkContext, seq: Seq[String]): RDD[String] = {
     sc.makeRDD(seq, 2).groupBy(x => x).flatMap(_._2)
   }