[SPARK-51955] Adding release() to ReadStateStore interface and reusing ReadStore for Streaming Aggregations

### What changes were proposed in this pull request?

Adding a release() method to the ReadStateStore interface to properly close read stores without aborting them
Implementing a getWriteStore() method that allows converting a read-only store to a writable store
Creating a StateStoreRDDProvider interface for tracking state stores by partition ID
Enhancing StateStoreRDD to find and reuse existing state stores through RDD lineage
Improving task completion handling with proper cleanup listeners

### Why are the changes needed?
Currently, stateful operations like aggregations follow a pattern where both read and write stores are opened simultaneously:
readStore.acquire()
writeStore.acquire()
writeStore.commit()
readStore.abort()
This pattern creates inefficiency because:

The abort() call on the read store unnecessarily invalidates the store's state, causing subsequent operations to reload the entire state store from scratch
Having two stores open simultaneously increases memory usage and can create contention issues
The upcoming lock hardening changes will only allow one state store to be open at a time, making this pattern incompatible

With the new approach, the usage paradigm becomes:
readStore = getReadStore()
writeStore = getWriteStore(readStore)
writeStore.commit()
This new paradigm allows us to reuse an existing read store by converting it to a write store using getWriteStore(), and properly clean up resources using release() instead of abort() when operations complete successfully. This avoids the unnecessary reloading of state data and improves performance while being compatible with future lock hardening changes.
### Does this PR introduce any user-facing change?
No
### How was this patch tested?
Unit tests
### Was this patch authored or co-authored using generative AI tooling?
No

Closes #51566 from ericm-db/read-store.

Authored-by: Eric Marnadi <[email protected]>
Signed-off-by: Anish Shrigondekar <[email protected]>

Author: ericm-db

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala

@@ -204,7 +204,6 @@ class StatePartitionReader(
   }
 
   override def close(): Unit = {
-    store.abort()
     super.close()
   }
 }

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

@@ -31,7 +31,7 @@ import org.apache.commons.io.IOUtils
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs._
 
-import org.apache.spark.{SparkConf, SparkEnv, SparkException}
+import org.apache.spark.{SparkConf, SparkEnv, SparkException, TaskContext}
 import org.apache.spark.internal.{Logging, LogKeys, MDC}
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
@@ -89,6 +89,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
 
     override def abort(): Unit = {}
 
+    override def release(): Unit = {}
+
     override def toString(): String = {
       s"HDFSReadStateStore[stateStoreId=$stateStoreId_, version=$version]"
     }
@@ -112,6 +114,15 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     case object UPDATING extends STATE
     case object COMMITTED extends STATE
     case object ABORTED extends STATE
+    case object RELEASED extends STATE
+
+    Option(TaskContext.get()).foreach { ctxt =>
+      ctxt.addTaskCompletionListener[Unit](ctx => {
+        if (state == UPDATING) {
+          abort()
+        }
+      })
+    }
 
     private val newVersion = version + 1
     @volatile private var state: STATE = UPDATING
@@ -953,7 +964,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
    * @param endVersion   checkpoint version to end with
    * @return [[HDFSBackedStateStore]]
    */
-  override def replayStateFromSnapshot(snapshotVersion: Long, endVersion: Long): StateStore = {
+  override def replayStateFromSnapshot(
+      snapshotVersion: Long, endVersion: Long, readOnly: Boolean): StateStore = {
     val newMap = replayLoadedMapFromSnapshot(snapshotVersion, endVersion)
     logInfo(log"Retrieved snapshot at version " +
       log"${MDC(LogKeys.STATE_STORE_VERSION, snapshotVersion)} and apply delta files to version " +

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

@@ -1025,6 +1025,12 @@ class RocksDB(
     }
   }
 
+  def release(): Unit = {
+    if (db != null) {
+      release(LoadStore)
+    }
+  }
+
   /**
    * Commit all the updates made as a version to DFS. The steps it needs to do to commits are:
    * - Flush all changes to disk

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

@@ -26,7 +26,7 @@ import scala.util.control.NonFatal
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
-import org.apache.spark.{SparkConf, SparkEnv}
+import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
 import org.apache.spark.internal.{Logging, MDC}
 import org.apache.spark.internal.LogKeys._
 import org.apache.spark.io.CompressionCodec
@@ -43,12 +43,15 @@ private[sql] class RocksDBStateStoreProvider
   with SupportsFineGrainedReplay {
   import RocksDBStateStoreProvider._
 
-  class RocksDBStateStore(lastVersion: Long) extends StateStore {
+  class RocksDBStateStore(
+      lastVersion: Long,
+      private[RocksDBStateStoreProvider] var readOnly: Boolean) extends StateStore {
     /** Trait and classes representing the internal state of the store */
     trait STATE
     case object UPDATING extends STATE
     case object COMMITTED extends STATE
     case object ABORTED extends STATE
+    case object RELEASED extends STATE
 
     @volatile private var state: STATE = UPDATING
     @volatile private var isValidated = false
@@ -57,6 +60,23 @@ private[sql] class RocksDBStateStoreProvider
 
     override def version: Long = lastVersion
 
+    Option(TaskContext.get()).foreach { ctxt =>
+      ctxt.addTaskCompletionListener[Unit](ctx => {
+        try {
+          if (state == UPDATING) {
+            if (readOnly) {
+              release()
+            } else {
+              abort() // Abort since this is an error if stateful task completes
+            }
+          }
+        } catch {
+          case NonFatal(e) =>
+            logWarning("Failed to abort state store", e)
+        }
+      })
+    }
+
     override def createColFamilyIfAbsent(
         colFamilyName: String,
         keySchema: StructType,
@@ -246,6 +266,19 @@ private[sql] class RocksDBStateStoreProvider
       }
     }
 
+    override def release(): Unit = {
+      assert(readOnly, "Release can only be called on a read-only store")
+      if (state != RELEASED) {
+        logInfo(log"Releasing ${MDC(VERSION_NUM, version + 1)} " +
+          log"for ${MDC(STATE_STORE_ID, id)}")
+        rocksDB.release()
+        state = RELEASED
+      } else {
+        // Optionally log at DEBUG level that it's already released
+        logDebug(log"State store already released")
+      }
+    }
+
     override def abort(): Unit = {
       verify(state == UPDATING || state == ABORTED, "Cannot abort after already committed")
       logInfo(log"Aborting ${MDC(VERSION_NUM, version + 1)} " +
@@ -450,38 +483,49 @@ private[sql] class RocksDBStateStoreProvider
 
   override protected def logName: String = s"${super.logName} ${stateStoreProviderId}"
 
-  override def getStore(version: Long, uniqueId: Option[String] = None): StateStore = {
+  /**
+   * Creates and returns a state store with the specified parameters.
+   *
+   * @param version The version of the state store to load
+   * @param uniqueId Optional unique identifier for checkpoint
+   * @param readOnly Whether to open the store in read-only mode
+   * @param existingStore Optional existing store to reuse instead of creating a new one
+   * @return The loaded state store
+   */
+  private def loadStateStore(
+      version: Long,
+      uniqueId: Option[String] = None,
+      readOnly: Boolean,
+      existingStore: Option[RocksDBStateStore] = None): StateStore = {
     try {
       if (version < 0) {
         throw QueryExecutionErrors.unexpectedStateStoreVersion(version)
       }
-      rocksDB.load(
-        version,
-        stateStoreCkptId = if (storeConf.enableStateStoreCheckpointIds) uniqueId else None)
-      new RocksDBStateStore(version)
-    }
-    catch {
-      case e: OutOfMemoryError =>
-        throw QueryExecutionErrors.notEnoughMemoryToLoadStore(
-          stateStoreId.toString,
-          "ROCKSDB_STORE_PROVIDER",
-          e)
-      case e: Throwable => throw StateStoreErrors.cannotLoadStore(e)
-    }
-  }
 
-  override def getReadStore(version: Long, uniqueId: Option[String] = None): StateStore = {
-    try {
-      if (version < 0) {
-        throw QueryExecutionErrors.unexpectedStateStoreVersion(version)
+      // Early validation of the existing store type before loading RocksDB
+      existingStore.foreach { store =>
+        if (!store.readOnly) {
+          throw new IllegalArgumentException(
+            s"Existing store must be readOnly, but got a read-write store")
+        }
       }
+
       rocksDB.load(
         version,
         stateStoreCkptId = if (storeConf.enableStateStoreCheckpointIds) uniqueId else None,
-        readOnly = true)
-      new RocksDBStateStore(version)
-    }
-    catch {
+        readOnly = readOnly)
+
+      // Create or reuse store instance
+      existingStore match {
+        case Some(store: RocksDBStateStore) =>
+          // Mark store as being used for write operations
+          store.readOnly = readOnly
+          store
+        case None =>
+          // Create new store instance
+          new RocksDBStateStore(version, readOnly)
+      }
+    } catch {
       case e: OutOfMemoryError =>
         throw QueryExecutionErrors.notEnoughMemoryToLoadStore(
           stateStoreId.toString,
@@ -491,6 +535,31 @@ private[sql] class RocksDBStateStoreProvider
     }
   }
 
+  override def getStore(
+      version: Long, uniqueId: Option[String] = None): StateStore = {
+    loadStateStore(version, uniqueId, readOnly = false)
+  }
+
+  override def upgradeReadStoreToWriteStore(
+      readStore: ReadStateStore,
+      version: Long,
+      uniqueId: Option[String] = None): StateStore = {
+    assert(version == readStore.version,
+      s"Can only upgrade readStore to writeStore with the same version," +
+        s" readStoreVersion: ${readStore.version}, writeStoreVersion: ${version}")
+    assert(this.stateStoreId == readStore.id, "Can only upgrade readStore to writeStore with" +
+      " the same stateStoreId")
+    assert(readStore.isInstanceOf[RocksDBStateStore], "Can only upgrade state store if it is a " +
+      "RocksDBStateStore")
+    loadStateStore(version, uniqueId, readOnly = false, existingStore =
+      Some(readStore.asInstanceOf[RocksDBStateStore]))
+  }
+
+  override def getReadStore(
+      version: Long, uniqueId: Option[String] = None): StateStore = {
+    loadStateStore(version, uniqueId, readOnly = true)
+  }
+
   override def doMaintenance(): Unit = {
     try {
       rocksDB.doMaintenance()
@@ -578,7 +647,8 @@ private[sql] class RocksDBStateStoreProvider
    * @param endVersion   checkpoint version to end with
    * @return [[StateStore]]
    */
-  override def replayStateFromSnapshot(snapshotVersion: Long, endVersion: Long): StateStore = {
+  override def replayStateFromSnapshot(
+      snapshotVersion: Long, endVersion: Long, readOnly: Boolean): StateStore = {
     try {
       if (snapshotVersion < 1) {
         throw QueryExecutionErrors.unexpectedStateStoreVersion(snapshotVersion)
@@ -587,7 +657,7 @@ private[sql] class RocksDBStateStoreProvider
         throw QueryExecutionErrors.unexpectedStateStoreVersion(endVersion)
       }
       rocksDB.loadFromSnapshot(snapshotVersion, endVersion)
-      new RocksDBStateStore(endVersion)
+      new RocksDBStateStore(endVersion, readOnly)
     }
     catch {
       case e: OutOfMemoryError =>