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

dylanwong250 · 2025-09-02T22:19:14Z

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:

[SPARK-53294][SS] Enable StateDataSource with state checkpoint v2 (only batchId option) #52047: Enable StateDataSource with state checkpoint v2 (only batchId option)
[SPARK-53333][SS] Enable StateDataSource with state checkpoint v2 (only readChangeFeed) #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

dylanwong250 · 2025-09-04T20:19:33Z

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

 /** Helper class for [[RocksDBCheckpointMetadata]] */
 object RocksDBCheckpointMetadata {
-  val VERSION = SQLConf.get.stateStoreCheckpointFormatVersion
+  val VERSION = 1


I want to point out this change. I believe using val VERSION = SQLConf.get.stateStoreCheckpointFormatVersion is incorrect in this case and I was seeing a few test failures. The QLConf.get.stateStoreCheckpointFormatVersion refers to the checkpoint format version not the metadata format version. Additionally, since this is inside an object as a val it will be instantiated once and is not the current value of SQLConf.get.stateStoreCheckpointFormatVersion. I think this VERSION may be shared incorrectly also.

I also experimented with having VERSION equal to the current value. I ran into a few issues with the maintenance threads having a different version than the streaming query they were ran in.

Thanks for catching this!

micheal-o · 2025-09-06T04:11:39Z

...he/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

+            .replayStateFromSnapshot(
+              opts.snapshotVersion,
+              opts.endVersion,
+              readOnly = false,


why is readOnly false for state reader?

I initially had it to readOnly = false to match what the replayStateFromSnapshot default value was without the uniqueIds. Also looking at the previous PR for snapshot reading #46944 it seems the readOnly option did not exist at the time. I changed to readOnly = true and all the tests pass so I think it is better to switch to readOnly = true.

...he/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

micheal-o · 2025-09-06T04:11:53Z

...he/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

+    }
+
+/** Snapshot options specialized for a single state store handler. */
+case class HandlerSnapshotOptions(


nit: private?

Made private[join]

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

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

liviazhu

LGTM, thanks!

liviazhu · 2025-09-08T21:23:13Z

...he/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

+    startKeyWithIndexToValueStateStoreCkptId: Option[String] = None,
+    endKeyToNumValuesStateStoreCkptId: Option[String] = None,
+    endKeyWithIndexToValueStateStoreCkptId: Option[String] = None) {
+      def getKeyToNumValuesHandlerOpts(): HandlerSnapshotOptions =


nit: indents seem wrong

Fixed. I used https://github.com/databricks/scala-style-guide?tab=readme-ov-file#spacing-and-indentation for reference.

anishshri-db · 2025-09-10T00:03:14Z

...g/apache/spark/sql/execution/datasources/v2/state/StreamStreamJoinStatePartitionReader.scala

-  private val keyToNumValuesStateStoreCkptId = if (joinSide == LeftSide) {
+  private val endStateStoreCheckpointIds =
+    SymmetricHashJoinStateManager.getStateStoreCheckpointIds(
+    partition.partition,


nit: indent for the args ?

anishshri-db · 2025-09-10T00:03:55Z

...he/spark/sql/execution/streaming/operators/stateful/join/SymmetricHashJoinStateManager.scala

-  protected class KeyWithIndexToValueStore(stateFormatVersion: Int)
-    extends StateStoreHandler(KeyWithIndexToValueType, keyWithIndexToValueStateStoreCkptId) {
+  protected class KeyWithIndexToValueStore(
+    stateFormatVersion: Int,


nit: this should be 4 spaces here ?

anishshri-db

lgtm pending nits

Add snapshotStartBatchId w checkpoint v2

365fe70

github-actions bot added SQL STRUCTURED STREAMING PYTHON labels Sep 2, 2025

Dylan Wong added 11 commits September 3, 2025 00:36

Add tests for StateDataSourceReadSuite

ed44243

remove comment

3484253

fix comments

9ab92e1

fix object val VERSION

0d80806

fix

f799147

Fix version coming from maintence thread

21c0070

fix tests

9767069

Checkpoint metadata version is 1

f908922

fix

db8c6fc

metadata version always 1

d1d32e1

Fix tests

efeb609

dylanwong250 commented Sep 4, 2025

View reviewed changes

micheal-o reviewed Sep 6, 2025

View reviewed changes

PR fixes

bbb800d

liviazhu approved these changes Sep 8, 2025

View reviewed changes

Fix formatting

a33b5fc

anishshri-db reviewed Sep 10, 2025

View reviewed changes

anishshri-db approved these changes Sep 10, 2025

View reviewed changes

Fix formatting

d35798e

anishshri-db closed this in feaf659 Sep 10, 2025

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

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

Uh oh!

Conversation

dylanwong250 commented Sep 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

dylanwong250 Sep 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

liviazhu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anishshri-db left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dylanwong250 commented Sep 2, 2025 •

edited

Loading

dylanwong250 Sep 4, 2025 •

edited

Loading