[SPARK-53332][SS] Enable StateDataSource with state checkpoint v2 (only snapshotStartBatchId option)

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

This PR enables StateDataSource support with state checkpoint v2 format for the `snapshotStartBatchId` and related options, completing the StateDataSource checkpoint v2 integration.

There is changes to the replayStateFromSnapshot method signature. `snapshotVersionStateStoreCkptId` and `endVersionStateStoreCkptId`. Both are needed as `snapshotVersionStateStoreCkptId` is used when getting the snapshot and `endVersionStateStoreCkptId` for calculating the full lineage from the final version.

Before
```
def replayStateFromSnapshot(
      snapshotVersion: Long, endVersion: Long, readOnly: Boolean = false): StateStore
```

After
```
def replayStateFromSnapshot(
      snapshotVersion: Long, endVersion: Long, readOnly: Boolean = false): StateStore
      snapshotVersion: Long,
      endVersion: Long,
      readOnly: Boolean = false,
      snapshotVersionStateStoreCkptId: Option[String] = None,
      endVersionStateStoreCkptId: Option[String] = None): StateStore
```

This is the final PR in the series following:
  - #52047: Enable StateDataSource with state checkpoint v2 (only batchId option)
  - #52148: Enable StateDataSource with state checkpoint v2 (only readChangeFeed)

NOTE: To read checkpoint v2 state data sources it is required to have `"spark.sql.streaming.stateStore.checkpointFormatVersion" -> 2`. It is possible to allow reading state data sources arbitrarily based on what is in the CommitLog by relaxing assertion checks but this is left as a future change.

### Why are the changes needed?

State checkpoint v2 (`"spark.sql.streaming.stateStore.checkpointFormatVersion"`) introduces a new format for storing state metadata that includes unique identifiers in the file path for each state store. The existing StateDataSource implementation only worked with checkpoint v1 format, making it incompatible with streaming queries using the newer checkpoint format. Only `batchId` was implemented in #52047 and only `readChangeFeed` was implemented in #52148.

### Does this PR introduce _any_ user-facing change?

Yes.

State Data Source will work when checkpoint v2 is used and the `snapshotStartBatchId` and related options are used.

### How was this patch tested?

In the previous PRs test suites were added to parameterize the current tests with checkpoint v2. All of these tests are now added back. All tests that previously intentionally tested some feature of the State Data Source Reader with checkpoint v1 should now be parameterized with checkpoint v2 (including python tests).

`RocksDBWithCheckpointV2StateDataSourceReaderSnapshotSuite` is added which uses the golden file approach similar to #46944 where `snapshotStartBatchId` is first added.

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #52202 from dylanwong250/SPARK-53332.

Authored-by: Dylan Wong <[email protected]>
Signed-off-by: Anish Shrigondekar <[email protected]>

Author: Dylan Wong

python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py

@@ -1923,10 +1923,6 @@ def conf(cls):
         cfg.set("spark.sql.streaming.stateStore.checkpointFormatVersion", "2")
         return cfg
 
-    # TODO(SPARK-53332): Add test back when checkpoint v2 support exists for snapshotStartBatchId
-    def test_transform_with_value_state_metadata(self):
-        pass
-
 
 class TransformWithStateInPySparkWithCheckpointV2TestsMixin(TransformWithStateInPySparkTestsMixin):
     @classmethod
@@ -1935,10 +1931,6 @@ def conf(cls):
         cfg.set("spark.sql.streaming.stateStore.checkpointFormatVersion", "2")
         return cfg
 
-    # TODO(SPARK-53332): Add test back when checkpoint v2 support exists for snapshotStartBatchId
-    def test_transform_with_value_state_metadata(self):
-        pass
-
 
 class TransformWithStateInPandasTests(TransformWithStateInPandasTestsMixin, ReusedSQLTestCase):
     pass

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

@@ -610,14 +610,6 @@ object StateSourceOptions extends DataSourceOptions {
       )
     }
 
-    if (startOperatorStateUniqueIds.isDefined) {
-      if (fromSnapshotOptions.isDefined) {
-        throw StateDataSourceErrors.invalidOptionValue(
-          SNAPSHOT_START_BATCH_ID,
-          "Snapshot reading is currently not supported with checkpoint v2.")
-      }
-    }
-
     StateSourceOptions(
       resolvedCpLocation, batchId.get, operatorId, storeName, joinSide,
       readChangeFeed, fromSnapshotOptions, readChangeFeedOptions,

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

@@ -206,7 +206,9 @@ class StatePartitionReader(
         provider.asInstanceOf[SupportsFineGrainedReplay]
           .replayReadStateFromSnapshot(
             fromSnapshotOptions.snapshotStartBatchId + 1,
-            partition.sourceOptions.batchId + 1)
+            partition.sourceOptions.batchId + 1,
+            getStartStoreUniqueId,
+            getEndStoreUniqueId)
     }
   }
 

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

@@ -24,7 +24,7 @@ import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, Par
 import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.JoinSideValues
 import org.apache.spark.sql.execution.datasources.v2.state.utils.SchemaUtil
 import org.apache.spark.sql.execution.streaming.operators.stateful.StatefulOperatorStateInfo
-import org.apache.spark.sql.execution.streaming.operators.stateful.join.{JoinStateManagerStoreGenerator, SymmetricHashJoinStateManager}
+import org.apache.spark.sql.execution.streaming.operators.stateful.join.{JoinStateManagerStoreGenerator, SnapshotOptions, SymmetricHashJoinStateManager}
 import org.apache.spark.sql.execution.streaming.operators.stateful.join.StreamingSymmetricHashJoinHelper.{JoinSide, LeftSide, RightSide}
 import org.apache.spark.sql.execution.streaming.state.StateStoreConf
 import org.apache.spark.sql.types.{BooleanType, StructType}
@@ -78,22 +78,40 @@ class StreamStreamJoinStatePartitionReader(
 
   private val startStateStoreCheckpointIds =
     SymmetricHashJoinStateManager.getStateStoreCheckpointIds(
-    partition.partition,
-    partition.sourceOptions.startOperatorStateUniqueIds,
-    usesVirtualColumnFamilies)
+      partition.partition,
+      partition.sourceOptions.startOperatorStateUniqueIds,
+      usesVirtualColumnFamilies)
 
-  private val keyToNumValuesStateStoreCkptId = if (joinSide == LeftSide) {
+  private val endStateStoreCheckpointIds =
+    SymmetricHashJoinStateManager.getStateStoreCheckpointIds(
+      partition.partition,
+      partition.sourceOptions.endOperatorStateUniqueIds,
+      usesVirtualColumnFamilies)
+
+  private val startKeyToNumValuesStateStoreCkptId = if (joinSide == LeftSide) {
     startStateStoreCheckpointIds.left.keyToNumValues
   } else {
     startStateStoreCheckpointIds.right.keyToNumValues
   }
 
-  private val keyWithIndexToValueStateStoreCkptId = if (joinSide == LeftSide) {
+  private val startKeyWithIndexToValueStateStoreCkptId = if (joinSide == LeftSide) {
     startStateStoreCheckpointIds.left.keyWithIndexToValue
   } else {
     startStateStoreCheckpointIds.right.keyWithIndexToValue
   }
 
+  private val endKeyToNumValuesStateStoreCkptId = if (joinSide == LeftSide) {
+    endStateStoreCheckpointIds.left.keyToNumValues
+  } else {
+    endStateStoreCheckpointIds.right.keyToNumValues
+  }
+
+  private val endKeyWithIndexToValueStateStoreCkptId = if (joinSide == LeftSide) {
+    endStateStoreCheckpointIds.left.keyWithIndexToValue
+  } else {
+    endStateStoreCheckpointIds.right.keyWithIndexToValue
+  }
+
   /*
    * This is to handle the difference of schema across state format versions. The major difference
    * is whether we have added new field(s) in addition to the fields from input schema.
@@ -150,13 +168,19 @@ class StreamStreamJoinStatePartitionReader(
         storeConf = storeConf,
         hadoopConf = hadoopConf.value,
         partitionId = partition.partition,
-        keyToNumValuesStateStoreCkptId = keyToNumValuesStateStoreCkptId,
-        keyWithIndexToValueStateStoreCkptId = keyWithIndexToValueStateStoreCkptId,
+        keyToNumValuesStateStoreCkptId = startKeyToNumValuesStateStoreCkptId,
+        keyWithIndexToValueStateStoreCkptId = startKeyWithIndexToValueStateStoreCkptId,
         formatVersion,
         skippedNullValueCount = None,
         useStateStoreCoordinator = false,
-        snapshotStartVersion =
-          partition.sourceOptions.fromSnapshotOptions.map(_.snapshotStartBatchId + 1),
+        snapshotOptions =
+          partition.sourceOptions.fromSnapshotOptions.map(opts => SnapshotOptions(
+            snapshotVersion = opts.snapshotStartBatchId + 1,
+            endVersion = partition.sourceOptions.batchId + 1,
+            startKeyToNumValuesStateStoreCkptId = startKeyToNumValuesStateStoreCkptId,
+            startKeyWithIndexToValueStateStoreCkptId = startKeyWithIndexToValueStateStoreCkptId,
+            endKeyToNumValuesStateStoreCkptId = endKeyToNumValuesStateStoreCkptId,
+            endKeyWithIndexToValueStateStoreCkptId = endKeyWithIndexToValueStateStoreCkptId)),
         joinStoreGenerator = new JoinStateManagerStoreGenerator()
       )
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

@@ -59,6 +59,8 @@ import org.apache.spark.util.NextIterator
  *                                  store providers being used in this class. If true, Spark will
  *                                  take care of management for state store providers, e.g. running
  *                                  maintenance task for these providers.
+ * @param snapshotOptions       Options controlling snapshot-based state replay for the state data
+ *                              source reader.
  * @param joinStoreGenerator    The generator to create state store instances, re-using the same
  *                              instance when the join implementation uses virtual column families
  *                              for join version 3.
@@ -95,15 +97,20 @@ abstract class SymmetricHashJoinStateManager(
     stateFormatVersion: Int,
     skippedNullValueCount: Option[SQLMetric] = None,
     useStateStoreCoordinator: Boolean = true,
-    snapshotStartVersion: Option[Long] = None,
+    snapshotOptions: Option[SnapshotOptions] = None,
     joinStoreGenerator: JoinStateManagerStoreGenerator) extends Logging {
   import SymmetricHashJoinStateManager._
 
   protected val keySchema = StructType(
     joinKeys.zipWithIndex.map { case (k, i) => StructField(s"field$i", k.dataType, k.nullable) })
   protected val keyAttributes = toAttributes(keySchema)
-  protected val keyToNumValues = new KeyToNumValuesStore(stateFormatVersion)
-  protected val keyWithIndexToValue = new KeyWithIndexToValueStore(stateFormatVersion)
+
+  protected val keyToNumValues = new KeyToNumValuesStore(
+    stateFormatVersion,
+    snapshotOptions.map(_.getKeyToNumValuesHandlerOpts()))
+  protected val keyWithIndexToValue = new KeyWithIndexToValueStore(
+    stateFormatVersion,
+    snapshotOptions.map(_.getKeyWithIndexToValueHandlerOpts()))
 
   /*
   =====================================================
@@ -456,7 +463,8 @@ abstract class SymmetricHashJoinStateManager(
   /** Helper trait for invoking common functionalities of a state store. */
   protected abstract class StateStoreHandler(
       stateStoreType: StateStoreType,
-      stateStoreCkptId: Option[String]) extends Logging {
+      stateStoreCkptId: Option[String],
+      handlerSnapshotOptions: Option[HandlerSnapshotOptions] = None) extends Logging {
     private var stateStoreProvider: StateStoreProvider = _
 
     /** StateStore that the subclasses of this class is going to operate on */
@@ -497,7 +505,7 @@ abstract class SymmetricHashJoinStateManager(
       }
       val storeProviderId = StateStoreProviderId(stateInfo.get, partitionId, storeName)
       val store = if (useStateStoreCoordinator) {
-        assert(snapshotStartVersion.isEmpty, "Should not use state store coordinator " +
+        assert(handlerSnapshotOptions.isEmpty, "Should not use state store coordinator " +
           "when reading state as data source.")
         joinStoreGenerator.getStore(
           storeProviderId, keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema),
@@ -509,13 +517,19 @@ abstract class SymmetricHashJoinStateManager(
           storeProviderId, keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema),
           useColumnFamilies = useVirtualColumnFamilies, storeConf, hadoopConf,
           useMultipleValuesPerKey = false, stateSchemaProvider = None)
-        if (snapshotStartVersion.isDefined) {
+        if (handlerSnapshotOptions.isDefined) {
           if (!stateStoreProvider.isInstanceOf[SupportsFineGrainedReplay]) {
             throw StateStoreErrors.stateStoreProviderDoesNotSupportFineGrainedReplay(
               stateStoreProvider.getClass.toString)
           }
+          val opts = handlerSnapshotOptions.get
           stateStoreProvider.asInstanceOf[SupportsFineGrainedReplay]
-            .replayStateFromSnapshot(snapshotStartVersion.get, stateInfo.get.storeVersion)
+            .replayStateFromSnapshot(
+              opts.snapshotVersion,
+              opts.endVersion,
+              readOnly = true,
+              opts.startStateStoreCkptId,
+              opts.endStateStoreCkptId)
         } else {
           stateStoreProvider.getStore(stateInfo.get.storeVersion, stateStoreCkptId)
         }
@@ -539,9 +553,12 @@ abstract class SymmetricHashJoinStateManager(
 
 
   /** A wrapper around a [[StateStore]] that stores [key -> number of values]. */
-  protected class KeyToNumValuesStore(val stateFormatVersion: Int)
-    extends StateStoreHandler(KeyToNumValuesType, keyToNumValuesStateStoreCkptId) {
-
+  protected class KeyToNumValuesStore(
+      val stateFormatVersion: Int,
+      val handlerSnapshotOptions: Option[HandlerSnapshotOptions] = None)
+    extends StateStoreHandler(
+      KeyToNumValuesType, keyToNumValuesStateStoreCkptId, handlerSnapshotOptions) {
+SnapshotOptions
     private val useVirtualColumnFamilies = stateFormatVersion == 3
     private val longValueSchema = new StructType().add("value", "long")
     private val longToUnsafeRow = UnsafeProjection.create(longValueSchema)
@@ -707,8 +724,11 @@ abstract class SymmetricHashJoinStateManager(
    * A wrapper around a [[StateStore]] that stores the mapping; the mapping depends on the
    * state format version - please refer implementations of [[KeyWithIndexToValueRowConverter]].
    */
-  protected class KeyWithIndexToValueStore(stateFormatVersion: Int)
-    extends StateStoreHandler(KeyWithIndexToValueType, keyWithIndexToValueStateStoreCkptId) {
+  protected class KeyWithIndexToValueStore(
+      stateFormatVersion: Int,
+      handlerSnapshotOptions: Option[HandlerSnapshotOptions] = None)
+    extends StateStoreHandler(
+      KeyWithIndexToValueType, keyWithIndexToValueStateStoreCkptId, handlerSnapshotOptions) {
 
     private val useVirtualColumnFamilies = stateFormatVersion == 3
     private val keyWithIndexExprs = keyAttributes :+ Literal(1L)
@@ -848,11 +868,11 @@ class SymmetricHashJoinStateManagerV1(
     stateFormatVersion: Int,
     skippedNullValueCount: Option[SQLMetric] = None,
     useStateStoreCoordinator: Boolean = true,
-    snapshotStartVersion: Option[Long] = None,
+    snapshotOptions: Option[SnapshotOptions] = None,
     joinStoreGenerator: JoinStateManagerStoreGenerator) extends SymmetricHashJoinStateManager(
   joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
   partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
-  stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotStartVersion,
+  stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
   joinStoreGenerator) {
 
   /** Commit all the changes to all the state stores */
@@ -927,11 +947,11 @@ class SymmetricHashJoinStateManagerV2(
     stateFormatVersion: Int,
     skippedNullValueCount: Option[SQLMetric] = None,
     useStateStoreCoordinator: Boolean = true,
-    snapshotStartVersion: Option[Long] = None,
+    snapshotOptions: Option[SnapshotOptions] = None,
     joinStoreGenerator: JoinStateManagerStoreGenerator) extends SymmetricHashJoinStateManager(
   joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
   partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
-  stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotStartVersion,
+  stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
   joinStoreGenerator) {
 
   /** Commit all the changes to the state store */
@@ -1034,20 +1054,20 @@ object SymmetricHashJoinStateManager {
       stateFormatVersion: Int,
       skippedNullValueCount: Option[SQLMetric] = None,
       useStateStoreCoordinator: Boolean = true,
-      snapshotStartVersion: Option[Long] = None,
+      snapshotOptions: Option[SnapshotOptions] = None,
       joinStoreGenerator: JoinStateManagerStoreGenerator): SymmetricHashJoinStateManager = {
     if (stateFormatVersion == 3) {
       new SymmetricHashJoinStateManagerV2(
         joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
         partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
-        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotStartVersion,
+        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
         joinStoreGenerator
       )
     } else {
       new SymmetricHashJoinStateManagerV1(
         joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
         partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
-        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotStartVersion,
+        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
         joinStoreGenerator
       )
     }
@@ -1280,3 +1300,36 @@ object SymmetricHashJoinStateManager {
     }
   }
 }
+
+/**
+ * Options controlling snapshot-based state replay for state data source reader.
+ */
+case class SnapshotOptions(
+    snapshotVersion: Long,
+    endVersion: Long,
+    startKeyToNumValuesStateStoreCkptId: Option[String] = None,
+    startKeyWithIndexToValueStateStoreCkptId: Option[String] = None,
+    endKeyToNumValuesStateStoreCkptId: Option[String] = None,
+    endKeyWithIndexToValueStateStoreCkptId: Option[String] = None) {
+
+  def getKeyToNumValuesHandlerOpts(): HandlerSnapshotOptions =
+    HandlerSnapshotOptions(
+      snapshotVersion = snapshotVersion,
+      endVersion = endVersion,
+      startStateStoreCkptId = startKeyToNumValuesStateStoreCkptId,
+      endStateStoreCkptId = endKeyToNumValuesStateStoreCkptId)
+
+  def getKeyWithIndexToValueHandlerOpts(): HandlerSnapshotOptions =
+    HandlerSnapshotOptions(
+      snapshotVersion = snapshotVersion,
+      endVersion = endVersion,
+      startStateStoreCkptId = startKeyWithIndexToValueStateStoreCkptId,
+      endStateStoreCkptId = endKeyWithIndexToValueStateStoreCkptId)
+}
+
+/** Snapshot options specialized for a single state store handler. */
+private[join] case class HandlerSnapshotOptions(
+    snapshotVersion: Long,
+    endVersion: Long,
+    startStateStoreCkptId: Option[String],
+    endStateStoreCkptId: Option[String])

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

@@ -972,10 +972,22 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
    *
    * @param snapshotVersion checkpoint version of the snapshot to start with
    * @param endVersion   checkpoint version to end with
+   * @param readOnly whether the state store should be read-only
+   * @param snapshotVersionStateStoreCkptId state store checkpoint ID of the snapshot version
+   * @param endVersionStateStoreCkptId state store checkpoint ID of the end version
    * @return [[HDFSBackedStateStore]]
    */
   override def replayStateFromSnapshot(
-      snapshotVersion: Long, endVersion: Long, readOnly: Boolean): StateStore = {
+      snapshotVersion: Long,
+      endVersion: Long,
+      readOnly: Boolean,
+      snapshotVersionStateStoreCkptId: Option[String] = None,
+      endVersionStateStoreCkptId: Option[String] = None): StateStore = {
+    if (snapshotVersionStateStoreCkptId.isDefined || endVersionStateStoreCkptId.isDefined) {
+      throw StateStoreErrors.stateStoreCheckpointIdsNotSupported(
+        "HDFSBackedStateStoreProvider does not support checkpointFormatVersion > 1 " +
+        "but a state store checkpointID is passed in")
+    }
     val newMap = replayLoadedMapFromSnapshot(snapshotVersion, endVersion)
     logInfo(log"Retrieved snapshot at version " +
       log"${MDC(LogKeys.STATE_STORE_VERSION, snapshotVersion)} and apply delta files to version " +
@@ -990,10 +1002,21 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
    *
    * @param snapshotVersion checkpoint version of the snapshot to start with
    * @param endVersion   checkpoint version to end with
+   * @param snapshotVersionStateStoreCkptId state store checkpoint ID of the snapshot version
+   * @param endVersionStateStoreCkptId state store checkpoint ID of the end version
    * @return [[HDFSBackedReadStateStore]]
    */
-  override def replayReadStateFromSnapshot(snapshotVersion: Long, endVersion: Long):
+  override def replayReadStateFromSnapshot(
+      snapshotVersion: Long,
+      endVersion: Long,
+      snapshotVersionStateStoreCkptId: Option[String] = None,
+      endVersionStateStoreCkptId: Option[String] = None):
     ReadStateStore = {
+    if (snapshotVersionStateStoreCkptId.isDefined || endVersionStateStoreCkptId.isDefined) {
+      throw StateStoreErrors.stateStoreCheckpointIdsNotSupported(
+        "HDFSBackedStateStoreProvider does not support checkpointFormatVersion > 1 " +
+        "but a state store checkpointID is passed in")
+    }
     val newMap = replayLoadedMapFromSnapshot(snapshotVersion, endVersion)
     logInfo(log"Retrieved snapshot at version " +
       log"${MDC(LogKeys.STATE_STORE_VERSION, snapshotVersion)} and apply delta files to version " +