What is the RocksDB state backend?

Summary

• RocksDB is a built-in state store provider for Spark Structured Streaming, added in Spark 3.2. It manages state in native memory and on local disk instead of the JVM heap, and Structured Streaming saves state changes to your checkpoint location.
• In open source Spark the HDFS-backed provider is still the default. Enable RocksDB by setting spark.sql.streaming.stateStore.providerClass to org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider.
• The Spark docs suggest it when a stateful query keeps millions of keys and JVM garbage collection pauses make micro-batch times vary.
• On Databricks, RocksDB is the default state store provider in Databricks Runtime 17.3 and above, and Databricks recommends changelog checkpointing for all stateful queries.

What is the RocksDB state backend?

In Spark, the "RocksDB state backend" is the RocksDB state store provider. The Spark docs define a state store as "a versioned key-value store which provides both read and write operations". Structured Streaming uses a state store provider to handle stateful operations across batches, and Spark ships two built-in providers: the HDFS state store provider and, as of Spark 3.2, the RocksDB state store provider. RocksDB itself is "an embeddable persistent key-value store for fast storage". The Spark provider uses it "to efficiently manage the state in the native memory and the local disk", and "any changes to this state are automatically saved by Structured Streaming to the checkpoint location you have provided".

How does RocksDB differ from the default HDFS provider?

	HDFS state store provider	RocksDB state store provider
Default	Yes, in open source Spark	No in open source Spark. Default in Databricks Runtime 17.3 and above
Where state lives	"in memory map in the first stage", then "backed by files in an HDFS-compatible file system". The state data is "maintained in the JVM memory of the executors"	"the native memory and the local disk", not the JVM heap
Fault tolerance	Updates are applied transactionally, and each set of updates increments the store's version	State changes are saved to the checkpoint location, with "full fault-tolerance guarantees (the same as default state management)"

When should I use RocksDB?

The docs name the case: stateful operations such as streaming aggregation, streaming dropDuplicates, stream-stream joins, mapGroupsWithState or flatMapGroupsWithState that "maintain millions of keys in the state" can hit "large JVM garbage collection (GC) pauses causing high variations in the micro-batch processing times". The default provider keeps state in executor JVM memory, so a large number of state objects puts memory pressure on the JVM.

How do I enable it?

from pyspark.sql import functions as F

# Set this before you start the query
spark.conf.set(
    "spark.sql.streaming.stateStore.providerClass",
    "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider")

events = spark.readStream.format("rate").option("rowsPerSecond", 10).load()

counts = (events
    .withWatermark("timestamp", "10 minutes")
    .groupBy(F.window("timestamp", "5 minutes"))
    .count())

query = (counts.writeStream
    .outputMode("update")
    .format("console")
    .option("checkpointLocation", "path/to/checkpoint/dir")
    .start())

The rate source is for testing and generates timestamp and value columns. The windowed count is a stateful aggregation, so its state goes to whichever provider is configured.
You cannot switch providers on an existing query. The docs say: "To read the previous state of the query properly, the class of state store provider should be unchanged." To change such configurations, you "discard the checkpoint and start a new query". The same note lists spark.sql.shuffle.partitions, because "state is partitioned via applying hash function to key".

What can I tune?

Config	Default	What it does
spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled	false	Uploads a changelog instead of a snapshot at commit
spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage	false	Bounds total RocksDB memory on a node (a soft limit)
spark.sql.streaming.stateStore.rocksdb.maxMemoryUsageMB	500	The memory limit when bounding is on
spark.sql.streaming.stateStore.rocksdb.trackTotalNumberOfRows	true	Tracks the row count. Turning it off can improve performance, but the count is then reported as 0
spark.sql.streaming.stateStore.rocksdb.compression	lz4	Compression type

The traditional mechanism uploads the manifest files and new SST files of each RocksDB instance. Changelog checkpointing uploads only the changes since the last checkpoint, and you must restart the query for the change to apply.

How does Databricks handle RocksDB?

• Default and class name. "RocksDB is the default state store provider in Databricks Runtime 17.3 and above." On earlier runtimes, set spark.sql.streaming.stateStore.providerClass to com.databricks.sql.streaming.state.RocksDBStateStoreProvider before starting the query.
• Changelog checkpointing. It is available in Databricks Runtime 13.3 LTS and above, Databricks recommends it for all stateful queries, and it is on by default in 17.3 and above.
• Memory. In Databricks Runtime 17.3 and above, Databricks caps RocksDB memory per node automatically. On earlier versions, set boundedMemoryUsage and maxMemoryUsageMB yourself.
• Pipelines. Databricks recommends RocksDB-based state management for stateful pipelines, and serverless pipelines manage state store configuration automatically.
• Metrics. Each state operator reports RocksDB metrics in customMetrics inside stateOperators in StreamingQueryProgress.

FAQs

Is RocksDB the default state store in Apache Spark?

No. The Spark docs call the HDFS backend provider the default, and RocksDB is opt-in. Databricks Runtime 17.3 and above changes that default.

How do I keep state from growing forever?

Use a watermark. It "lets the engine automatically track the current event time in the data and attempt to clean up old state accordingly". For aggregations, the output mode must be Append or Update, and withWatermark must use the aggregation's timestamp column and come before the aggregation.

Can I inspect the state?

Yes. The State Data Source is experimental and, as of Spark 4.0, read-only: spark.read.format("statestore").load("<checkpointLocation>").