Okay this might be the way to go

Author: ericm-db

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

@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.streaming.state
 import java.io._
 import java.util.UUID
 import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
 
 import scala.util.control.NonFatal
 
@@ -58,14 +59,14 @@ private[sql] class RocksDBStateStoreProvider
 
     @volatile private var state: STATE = UPDATING
 
-    @volatile private var usedToCreateWriteStore: Boolean = false
+    private val usedToCreateWriteStore: AtomicBoolean = new AtomicBoolean(false)
 
     override def getReadStamp: Long = {
-      usedToCreateWriteStore = true
+      usedToCreateWriteStore.set(true)
       stamp
     }
 
-    override def usedForWriteStore: Boolean = usedToCreateWriteStore
+    override def usedForWriteStore: Boolean = usedToCreateWriteStore.get()
 
     /**
      * Validates the expected state, throws exception if state is not as expected.
@@ -132,7 +133,9 @@ private[sql] class RocksDBStateStoreProvider
     Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit] {
       _ =>
         try {
-          abort()
+          if (state == UPDATING) {
+            abort()
+          }
         } catch {
           case NonFatal(e) =>
             logWarning("Failed to abort state store", e)
@@ -337,7 +340,7 @@ private[sql] class RocksDBStateStoreProvider
     private var storedMetrics: Option[RocksDBMetrics] = None
 
     override def commit(): Long = synchronized {
-      if (usedToCreateWriteStore) {
+      if (usedToCreateWriteStore.get()) {
         return -1
       }
       validateState(List(UPDATING))
@@ -359,7 +362,7 @@ private[sql] class RocksDBStateStoreProvider
     }
 
     override def abort(): Unit = {
-      if (usedToCreateWriteStore) {
+      if (usedToCreateWriteStore.get()) {
         return
       }
       if (validateState(List(UPDATING, ABORTED)) != ABORTED) {

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

@@ -66,7 +66,7 @@ object StateStoreEncoding {
  *   not supported yet from the implementation. Note that some stateful operations would not work
  *   on disabling prefixScan functionality.
  */
-trait ReadStateStore {
+trait ReadStateStore extends Logging {
 
   /** Unique identifier of the store */
   def id: StateStoreId
@@ -937,6 +937,30 @@ object StateStore extends Logging {
       stateSchemaBroadcast)
     storeProvider.getStore(version, stateStoreCkptId)
   }
+
+  def getWriteStore(
+      readStore: ReadStateStore,
+      storeProviderId: StateStoreProviderId,
+      keySchema: StructType,
+      valueSchema: StructType,
+      keyStateEncoderSpec: KeyStateEncoderSpec,
+      version: Long,
+      stateStoreCkptId: Option[String],
+      stateSchemaBroadcast: Option[StateSchemaBroadcast],
+      useColumnFamilies: Boolean,
+      storeConf: StateStoreConf,
+      hadoopConf: Configuration,
+      useMultipleValuesPerKey: Boolean = false): StateStore = {
+    hadoopConf.set(StreamExecution.RUN_ID_KEY, storeProviderId.queryRunId.toString)
+    if (version < 0) {
+      throw QueryExecutionErrors.unexpectedStateStoreVersion(version)
+    }
+    hadoopConf.set(StreamExecution.RUN_ID_KEY, storeProviderId.queryRunId.toString)
+    val storeProvider = getStateStoreProvider(storeProviderId, keySchema, valueSchema,
+      keyStateEncoderSpec, useColumnFamilies, storeConf, hadoopConf, useMultipleValuesPerKey,
+      stateSchemaBroadcast)
+    storeProvider.getWriteStore(readStore, version, stateStoreCkptId)
+  }
   // scalastyle:on
 
   private def getStateStoreProvider(

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

@@ -27,6 +27,10 @@ import org.apache.spark.sql.internal.SessionState
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.SerializableConfiguration
 
+trait StateStoreRDDProvider {
+  def getStateStoreForPartition(partitionId: Int): Option[ReadStateStore]
+}
+
 abstract class BaseStateStoreRDD[T: ClassTag, U: ClassTag](
     dataRDD: RDD[T],
     checkpointLocation: String,
@@ -82,7 +86,17 @@ class ReadStateStoreRDD[T: ClassTag, U: ClassTag](
     useColumnFamilies: Boolean = false,
     extraOptions: Map[String, String] = Map.empty)
   extends BaseStateStoreRDD[T, U](dataRDD, checkpointLocation, queryRunId, operatorId,
-    sessionState, storeCoordinator, extraOptions) {
+    sessionState, storeCoordinator, extraOptions) with StateStoreRDDProvider {
+
+  // ThreadLocal to store state stores by partition ID
+  @transient private lazy val partitionStores =
+  new ThreadLocal[Map[Int, ReadStateStore]]() {
+    override def initialValue(): Map[Int, ReadStateStore] = Map.empty
+  }
+
+  override def getStateStoreForPartition(partitionId: Int): Option[ReadStateStore] = {
+    Option(partitionStores.get()).flatMap(_.get(partitionId))
+  }
 
   override protected def getPartitions: Array[Partition] = dataRDD.partitions
 
@@ -95,6 +109,8 @@ class ReadStateStoreRDD[T: ClassTag, U: ClassTag](
       stateStoreCkptIds.map(_.apply(partition.index).head),
       stateSchemaBroadcast,
       useColumnFamilies, storeConf, hadoopConfBroadcast.value.value)
+    // Store reference for this partition
+    partitionStores.set(partitionStores.get() + (partition.index -> store))
     storeReadFunction(store, inputIter)
   }
 }
@@ -126,16 +142,59 @@ class StateStoreRDD[T: ClassTag, U: ClassTag](
 
   override protected def getPartitions: Array[Partition] = dataRDD.partitions
 
+  // Recursively find a state store provider in the RDD lineage
+  private def findStateStoreProvider(rdd: RDD[_]): Option[StateStoreRDDProvider] = {
+    rdd match {
+      case provider: StateStoreRDDProvider => Some(provider)
+      case _ if rdd.dependencies.isEmpty => None
+      case _ =>
+        // Search all dependencies
+        rdd.dependencies.view
+          .map(dep => findStateStoreProvider(dep.rdd))
+          .find(_.isDefined)
+          .flatten
+    }
+  }
+
   override def compute(partition: Partition, ctxt: TaskContext): Iterator[U] = {
     val storeProviderId = getStateProviderId(partition)
-
     val inputIter = dataRDD.iterator(partition, ctxt)
-    val store = StateStore.get(
-      storeProviderId, keySchema, valueSchema, keyStateEncoderSpec, storeVersion,
-      uniqueId.map(_.apply(partition.index).head),
-      stateSchemaBroadcast,
-      useColumnFamilies, storeConf, hadoopConfBroadcast.value.value,
-      useMultipleValuesPerKey)
+
+    // Try to find a state store provider in the RDD lineage
+    val store = findStateStoreProvider(dataRDD).flatMap { provider =>
+      provider.getStateStoreForPartition(partition.index)
+    } match {
+      case Some(readStore) =>
+        // Convert the read store to a writable store
+        StateStore.getWriteStore(
+          readStore,
+          storeProviderId,
+          keySchema,
+          valueSchema,
+          keyStateEncoderSpec,
+          storeVersion,
+          uniqueId.map(_.apply(partition.index).head),
+          stateSchemaBroadcast,
+          useColumnFamilies,
+          storeConf,
+          hadoopConfBroadcast.value.value,
+          useMultipleValuesPerKey)
+
+      case None =>
+        // Fall back to creating a new store
+        StateStore.get(
+          storeProviderId,
+          keySchema,
+          valueSchema,
+          keyStateEncoderSpec,
+          storeVersion,
+          uniqueId.map(_.apply(partition.index).head),
+          stateSchemaBroadcast,
+          useColumnFamilies,
+          storeConf,
+          hadoopConfBroadcast.value.value,
+          useMultipleValuesPerKey)
+    }
     storeUpdateFunction(store, inputIter)
   }
 }