What are the state-store tuning best practices?

Summary

• State stores hold the intermediate data that stateful operators (aggregations, deduplication, stream-stream joins, and arbitrary stateful operators) keep between micro-batches. Tuning comes down to partition count, the state store provider, memory limits, and watermarks.
• spark.sql.shuffle.partitions (default 200) sets how many partitions state is split into, and it cannot be changed after a query has run without discarding the checkpoint.
• Spark has two built-in state store providers: the default HDFS-backed provider, and a RocksDB provider (since Spark 3.2) that keeps state in native memory and local disk instead of the JVM heap.
• Watermarks bound state: without one, deduplication stores the data from all past records and aggregation state is never dropped.
• On Databricks, RocksDB is the default provider from Databricks Runtime 17.3, and Databricks recommends changelog checkpointing for every stateful query.

What are the state-store tuning best practices?

In Apache Spark Structured Streaming, a state store is a versioned key-value store that stateful operations use across micro-batches. Stateful operations include streaming aggregation, streaming dropDuplicates, stream-stream joins, mapGroupsWithState, and flatMapGroupsWithState. The practices below follow the Apache Spark guide first, then the Databricks additions.

How do I choose the partition count?

State is partitioned by applying a hash function to the key, so the number of state partitions has to stay the same for the life of a checkpoint. Spark lists three configurations that cannot be modified after a query has run: spark.sql.shuffle.partitions (default 200), spark.sql.streaming.stateStore.providerClass, and spark.sql.streaming.multipleWatermarkPolicy. To change any of them, discard the checkpoint and start a new query.
To run fewer tasks for stateful operations, the guide points to coalesce, which avoids unnecessary repartitioning. After coalesce, the reduced number of tasks is kept unless another shuffle happens.

Which state store provider should I use?

Provider	Where state lives	Documented trade-off
HDFS-backed (default)	In a memory map in the executor JVM, backed by files in an HDFS-compatible file system	A large number of state objects puts memory pressure on the JVM and can cause long garbage collection (GC) pauses
RocksDB (Spark 3.2 and later)	Native memory and local disk, with changes saved to the checkpoint location	Avoids holding state in JVM memory, with the same fault-tolerance guarantees as the default

The Spark guide suggests RocksDB when a stateful query maintains millions of keys. To enable it in open-source Spark, set the provider class before the query first runs:

spark.conf.set(
    "spark.sql.streaming.stateStore.providerClass",
    "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider")

How do I tune RocksDB memory and checkpointing?

These settings and defaults come from the Spark guide.

Setting	Default	Effect
spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled	false	Uploads only the changes since the last checkpoint, not full snapshots
spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage	false	Caps total RocksDB memory on a node (a soft limit)
spark.sql.streaming.stateStore.rocksdb.maxMemoryUsageMB	500	The cap that applies when bounded memory is on
spark.sql.streaming.stateStore.rocksdb.trackTotalNumberOfRows	true	Turning it off can improve performance, but the total row count then reports 0

You can switch changelog checkpointing on or off between restarts without discarding existing state, but the query must restart for the change to apply.

How do watermarks keep state small?

For a window ending at time T, the engine keeps state until the maximum event time seen minus the late threshold passes T, then drops it. Complete output mode cannot drop state, because it preserves all aggregate data. Streaming deduplication without a watermark stores the data from all past records. Call withWatermark on the same column the aggregation uses, and before the aggregation:

# words: streaming DataFrame with columns timestamp and word
windowedCounts = words \
    .withWatermark("timestamp", "10 minutes") \
    .groupBy(
        window(words.timestamp, "10 minutes", "5 minutes"),
        words.word) \
    .count()

A larger threshold waits longer for late data but keeps more state, which might need more compute.

What should I watch while the query runs?

Spark runs each state store provider on the same executor across batches when it can, because reloading a large state from the checkpoint can be slow. Raising spark.locality.wait makes Spark wait longer for a data-local slot. With the HDFS-backed provider, watch the loadedMapCacheHitCount and loadedMapCacheMissCount metrics.
Each progress report has a stateOperators array with numRowsTotal, numRowsUpdated, numRowsRemoved, memoryUsedBytes, numRowsDroppedByWatermark, commitTimeMs, and numShufflePartitions. The Spark web UI's Structured Streaming tab shows aggregated state rows and state memory. Spark's State Data Source (experimental, read-only as of Spark 4.0) lets a batch query read the state in a checkpoint.

How does Databricks approach state store tuning?

• Provider: 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.
• Changelog checkpointing: Databricks recommends it for all stateful queries (Databricks Runtime 13.3 LTS and above). It is on by default from 17.3.
• Memory: From Databricks Runtime 17.3, Databricks caps RocksDB memory per node automatically. Earlier runtimes need the bounded-memory settings above.
• Cluster and partitions: Databricks recommends compute-optimized workers and shuffle partitions at 1-2 times the number of cores. The count is fixed when the checkpoint is created.
• Micro-batches: Databricks recommends setting spark.sql.streaming.noDataMicroBatches.enabled to false. Watermark-based output may then wait for new data.
• Resizing: On Databricks Runtime 18 LTS and above, on-demand state repartitioning (Public Preview) changes the partition count of a RocksDB-backed query through spark.sql.streaming.stateStore.partitions, without a new checkpoint.

Serverless pipelines manage state store configurations automatically, and state rebalancing is on by default for Lakeflow Pipelines.

FAQs

Can I change spark.sql.shuffle.partitions mid-stream?

Not in open-source Spark: state is hash-partitioned by key, so you must discard the checkpoint and start a new query. On Databricks, changing it has no effect on a query that already has a checkpoint, unless you use on-demand state repartitioning (Public Preview, Databricks Runtime 18 LTS and above).

Is RocksDB the only state store backend?

No. Spark has two built-in providers (HDFS-backed by default, and RocksDB), and you can write your own by extending the StateStoreProvider interface.