Spark Structured Streaming RocksDB state store: slow commits, executor memory blowups and checkpoint storage growth

The query already uses RocksDBStateStoreProvider, and its watermarks are correct, so state is no longer growing without limit. Batches still take minutes when they used to take seconds, executors get killed for exceeding memory, or the checkpoint folder is in the terabytes. This post works through those problems from the metrics Spark already reports. If state itself grows every day, fix watermarks and deduplication first. Tuning RocksDB will not save a query that never drops state.

#How the state store spends a batch

Each stateful operator has one state store instance per shuffle partition. In a micro-batch, each task:

1. Loads the store for that partition and version. If the executor already has it cached, this is quick. If not, the store is downloaded and rebuilt from the checkpoint.
2. Reads and writes keys (get and put).
3. Commits. It flushes memtables, optionally compacts, creates a local checkpoint, then uploads to the checkpoint location. That means changed SST files, or only a changelog file when changelog checkpointing is on.

Background maintenance then takes snapshots and deletes old versions. Slow batches almost always come down to step 1 (reloads) or step 3 (commit and upload).

Warning

The number of state store instances is fixed by spark.sql.shuffle.partitions on the query's first run, and it is stored in the checkpoint. Changing it later has no effect on an existing checkpoint. If 200 partitions are too few for your state size, the only fix is a new checkpoint and a state rebuild.

#Step 1: read the commit-phase metrics

import json
p = query.lastProgress
for op in p["stateOperators"]:
    print(op["operatorName"], op["numRowsTotal"], op["memoryUsedBytes"])
    print(json.dumps({k: v for k, v in op["customMetrics"].items()
                      if k.startswith("rocksdb")}, indent=1))
print(p["durationMs"])

The customMetrics worth reading are summed across tasks. Names vary by version, so check your build's output.

Also compare durationMs.addBatch with durationMs.commitOffsets and walCommit. If addBatch dominates but the commit metrics are small, the cost is in loading the store or in your own processing. Check the Spark UI for tasks with a long scheduler delay or deserialization time on the stateful stage.

#Symptom → cause → fix

#Fix 1: changelog checkpointing

With this off, each commit makes RocksDB produce a checkpoint, and every changed SST file is uploaded. With it on, the commit uploads a small changelog of the batch's changes, and full snapshots happen in the background.

spark.conf.set(
  "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled", "true")

Existing checkpoints can switch to it on restart, and newer runtimes can switch back. Databricks turns it on by default from Runtime 17.3; check your version's notes. The trade-off is recovery time. After a restart, a task replays changelogs since the last snapshot, so spark.sql.streaming.stateStore.minDeltasForSnapshot (default 10) caps how many changelogs it replays.

#Fix 2: bound off-heap memory

Each RocksDB instance has its own memtables and block cache, and none of it is JVM heap. An executor running 8 tasks may host many instances: one per partition per stateful operator, plus instances it keeps for reuse. Bound them all with a shared limit:

spark.conf.set("spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage", "true")
spark.conf.set("spark.sql.streaming.stateStore.rocksdb.maxMemoryUsageMB", "4096")
# Optional: share of that budget used for memtables and for high-priority (index/filter) blocks
spark.conf.set("spark.sql.streaming.stateStore.rocksdb.writeBufferCacheRatio", "0.5")
spark.conf.set("spark.sql.streaming.stateStore.rocksdb.highPriorityPoolRatio", "0.1")

Then raise spark.executor.memoryOverhead to at least the bound plus normal overhead. Otherwise the container limit, not RocksDB, decides when the executor dies. With bounded memory on, per-instance memory metrics may read 0, because usage is accounted at the shared cache. Watch container memory from the cluster or node metrics instead.

Tip

A cap that is too small turns memory problems into latency problems: rocksdbGetLatency and cache misses go up. Size the cap from rocksdbSstFileSize divided by executors. A working set of hot keys that fits in the cache is what you want.

#Fix 3: stop reloading state

A store instance stays loaded on the executor that last ran its partition. Spark tries to schedule each partition's task on that executor again. Losing the executor means the next batch downloads and rebuilds that partition's state from the checkpoint. With large state, that can take minutes per task.

• Run stateful queries on fixed-size compute. Turn off autoscaling and dynamic allocation for these clusters, and keep stateful stages off spot capacity.
• Don't share the cluster with batch jobs that churn executors.
• If you see reloads with no executor loss, check spark.locality.wait. A value of 0 lets tasks move away from the executor holding their state.

#Fix 4: checkpoint storage growth

spark.sql.streaming.minBatchesToRetain = 100   (default; versions of state and metadata kept)
spark.sql.streaming.stateStore.maintenanceInterval = 60s

Lower minBatchesToRetain if you never roll back more than a few batches. It bounds how far back you can restart, so don't go below what your recovery runbook needs. If the folder still grows, check that maintenance actually runs. It runs on executors in the background, and very short-lived clusters may never get to it. Never delete files under state/ by hand while the query is running, or between runs. A missing file makes the next restart fail.

#Inspect the state directly

Spark 4.0 and Databricks Runtime 14.3+ can read a checkpoint's state as a table. This is useful for finding keys that never expire:

meta = spark.read.format("state-metadata").load("/chk/orders_dedup")
meta.show(truncate=False)

state = spark.read.format("statestore").load("/chk/orders_dedup")
state.selectExpr("partition_id").groupBy("partition_id").count() \
     .orderBy("count", ascending=False).show(10)

A few partitions with far more rows than the rest means key skew in the state. RocksDB tuning won't help with that. Change the grouping key, or rebuild with more partitions.

Caution

Changing the state store provider, the stateful operator chain, or the shuffle partition count on an existing checkpoint is not supported, and the query can fail on restart. Schedule those changes as a new checkpoint with a planned replay, as you would any other checkpoint reset.

#Prevent it next time

• Start every stateful query with RocksDB, changelog checkpointing and bounded memory on, and choose spark.sql.shuffle.partitions for the state size you expect in a year.
• Ship stateOperators.customMetrics to your metrics store with a StreamingQueryListener. Alert on rocksdbCommitFileSyncLatencyMs and on batch duration trends, not only on failures.
• Keep stateful streams on dedicated, fixed-size compute.