A streaming query can work well when it first goes into production and become slow a few months later. The amount of incoming data may increase, the number of active users may grow, or the query may start tracking more records in memory. At that point, developers often need to change how Spark divides the work across partitions.
With traditional stateful Structured Streaming, changing the partition count is difficult because the existing checkpoint records the partition layout used by the query. Changing spark.sql.shuffle.partitions does not change that saved layout. Developers who need a different partition count have generally had to start with a new checkpoint, which can mean losing the saved streaming state.
Databricks has introduced on-demand state repartitioning to address this problem. The feature lets developers increase or decrease the number of state partitions while keeping the existing checkpoint.
The feature is in Public Preview and requires Databricks Runtime 18 LTS or later, with the RocksDB state store. It is useful for teams whose streaming workloads change over time and who want to adjust capacity without rebuilding their saved state.
Why Partitioning Matters in Stateful Streaming
Apache Spark processes large datasets by dividing them into partitions. Each partition can be processed as a separate unit of work, allowing multiple tasks to run across the available compute resources.
Stateful streaming adds another requirement. Spark must remember information between micro-batches so that it can calculate results correctly as new records arrive.
Consider a streaming application that counts events for each device. The application receives records from thousands of devices and maintains a running count for each device ID.
from pyspark.sql import functions as F
events = (
spark.readStream
.format("delta")
.load("/data/device-events")
)
device_counts = (
events
.groupBy("device_id")
.count()
)
The groupBy() operation requires Spark to maintain state for the device IDs it has encountered. When new events arrive, Spark updates the existing counts instead of treating every micro-batch as an unrelated dataset.
Other stateful operations include streaming aggregations, deduplication, stream-to-stream joins, and custom stateful processing.
The number of partitions affects how this work is distributed. Too few partitions can leave available CPU resources underused and create large, slow-running tasks. Too many partitions can add scheduling and state-management overhead.
The right number depends on the workload, the size of the saved state, and the compute resources available.
Why Changing Partitions Was Difficult
In a typical stateful Structured Streaming application, the checkpoint stores information about processed source records, committed micro-batches, and the state maintained by stateful operators.
When the query starts for the first time, Spark establishes the partition layout associated with the checkpoint. Later, changing spark.sql.shuffle.partitions alone does not replace that saved layout.
For example, suppose a query started with 200 partitions. The team later increases the cluster size and changes the configuration to 400 partitions.
spark.conf.set("spark.sql.shuffle.partitions", "400")
For an existing stateful query, this setting alone does not change the partition count recorded in its checkpoint. The query continues using its existing layout.
Previously, applying a new partition count generally required a new checkpoint. Starting with a new checkpoint means the query no longer has the old checkpoint's saved progress and state available for recovery. Depending on the source and sink, rebuilding the state may require replaying historical data or running a separate recovery process.
On-demand state repartitioning provides another option. Developers can change the partition count and restart the query using the same checkpoint location.
How On-Demand State Repartitioning Works
The feature uses a new Spark configuration:
spark.conf.set(
"spark.sql.streaming.stateStore.partitions",
"400"
)
This setting controls the number of partitions used by the state store. For stateful queries, it takes precedence over spark.sql.shuffle.partitions.
The change does not take effect simply because the configuration value has been updated. You must stop and restart the streaming query.
During the restart, Databricks completes the last planned micro-batch and then performs a repartition operation. That operation redistributes the saved state into the new number of partitions. Once repartitioning finishes, the query resumes normal processing.
The checkpoint remains in use throughout this process. The saved state is redistributed instead of being discarded.
There is still a period when the query is not processing new micro-batches normally, so this feature should not be confused with a zero-downtime scaling mechanism.
Practical Example: Changing from 200 to 400 Partitions
Suppose an application reads device events from a Delta table and calculates counts for each device over five-minute event-time windows.
The original query uses 200 state partitions. As the workload grows, you decide to increase the count to 400.
The following example shows how to start the query, stop it, update the configuration, and restart it with the same checkpoint.
from pyspark.sql import functions as F
source_path = "/data/device-events"
checkpoint_path = "/checkpoints/device-counts"
target_path = "/data/device-count-results"
def start_stream():
events = (
spark.readStream
.format("delta")
.load(source_path)
)
counts = (
events
.withWatermark("event_time", "10 minutes")
.groupBy(
F.window("event_time", "5 minutes"),
"device_id"
)
.count()
)
return (
counts.writeStream
.format("delta")
.outputMode("append")
.option("checkpointLocation", checkpoint_path)
.start(target_path)
)
# Start with 200 partitions.
spark.conf.set(
"spark.sql.streaming.stateStore.partitions",
"200"
)
query = start_stream()
# Stop the query before changing its partition count.
query.stop()
# Request 400 state partitions.
spark.conf.set(
"spark.sql.streaming.stateStore.partitions",
"400"
)
# Restart using the same checkpoint location.
query = start_stream()
This example assumes that the source contains an event_time column, that the target supports the selected output mode, and that the checkpoint and target paths are configured correctly for the environment.
The important detail is that checkpoint_path stays the same. Do not delete the checkpoint or generate a new path just to change the partition count. Doing so would defeat the purpose of retaining the saved state.
The example uses the same query definition before and after the restart. Changing the query's stateful operations or their state schema is a separate concern and is not made safe simply by enabling repartitioning.
When Should You Increase or Decrease the Partition Count?
Changing the partition count can help when a streaming workload no longer fits its original resource allocation. It is not a setting that should be increased automatically whenever processing slows down.
Situation | Possible action | What to check |
|---|
Micro-batches take longer as data volume grows | Increase partitions and measure the result | CPU utilization, task duration, state size |
A large number of records creates uneven task durations | Test a different partition count | Data skew and the distribution of state keys |
The cluster has many idle cores | Consider reducing partitions or reviewing the compute size | Scheduling overhead and task parallelism |
State repartitioning takes a long time | Plan a suitable maintenance window | State size and repartition duration |
The query is slow because of a source or sink bottleneck | Investigate the source or sink first | Input delays, output latency, and external service limits |
Increasing the number of partitions does not guarantee faster processing. If a database sink is the bottleneck, adding state partitions may have little effect. If a few keys contain most of the data, simply increasing the partition count might not solve the underlying skew.
Measure the query before and after the change to determine whether the adjustment helped.
Requirements and Limitations
On-demand state repartitioning has specific requirements.
Databricks Runtime: Use Databricks Runtime 18 LTS or later.
State store: The query must use the RocksDB state store provider. RocksDB is the default provider in Databricks Runtime 17.3 and later.
Checkpoint: Stop and restart the query with its existing checkpoint location.
Restart behavior: The query performs a repartition operation after the last planned micro-batch completes.
Availability: The feature is in Public Preview, so check its current limitations before relying on it for a critical workload.
A larger state can take longer to redistribute. Plan the change around the application's latency requirements, and do not assume that the query will resume immediately after the restart.
If you are using an older runtime, verify whether upgrading is appropriate for your workload. Do not enable the feature on an unsupported runtime and assume that changing the configuration will make it work.
How to Monitor Repartitioning
After the query restarts, monitor its progress and compare the processing duration with the previous runs.
Databricks exposes the repartition operation through StreamingQueryProgress events. The durationMs metrics include controlBatch.REPARTITION, which reports the time spent on the repartition operation in milliseconds.
For example, a monitoring system can record the duration and alert when it exceeds an operational threshold.
def report_repartition_duration(progress):
durations = progress.durationMs or {}
repartition_ms = durations.get("controlBatch.REPARTITION")
if repartition_ms is not None:
print(
f"State repartition completed in "
f"{repartition_ms} ms"
)
This is a small illustration of reading the progress event, not a complete streaming listener. In a production monitoring setup, connect the handler to the query's progress events and send the measurement to your logging or metrics system.
Record the duration before and after changes to the partition count. Also track normal micro-batch processing time, input volume, and state size. These measurements help you distinguish a slow repartition operation from a query that remains slow after repartitioning.
Advantages and Disadvantages
Advantages
Keeps checkpoint state: Developers can change the partition count without creating a new checkpoint.
Supports changing workloads: Teams can adjust partitioning as input volume and compute capacity change.
Reduces recovery work: The existing checkpoint retains the query's saved progress and state.
Allows scaling in either direction: You can increase or decrease the number of partitions based on measured needs.
Disadvantages
Requires a supported runtime: Older Databricks Runtime versions do not meet the feature's requirements.
Needs a query restart: This is not an instant configuration change.
Repartitioning takes time: Larger state can increase the duration of the operation.
Preview limitations apply: The feature is not yet generally available, so teams should test it carefully.
Best Practices for Production Workloads
Measure before changing the configuration. Use query progress metrics to identify whether the partition count is actually causing the problem.
Keep the checkpoint safe. Use durable storage for checkpoint files and never delete the checkpoint as part of a routine repartitioning procedure.
Change one variable at a time. If you increase partitions and resize the cluster simultaneously, it becomes harder to determine which change improved or worsened performance.
Test with representative state. A small development query may repartition quickly, while a production query with a large state store may take considerably longer.
Plan for the restart. Choose a suitable window, communicate the expected processing pause, and confirm that downstream consumers can handle the delay.
Monitor the next few micro-batches. A successful repartition does not guarantee that the new configuration is optimal. Compare task duration, processing latency, and resource utilization after the query resumes.
Conclusion
On-demand state repartitioning gives Databricks developers a way to adjust stateful streaming workloads without discarding their existing checkpoint state. By changing spark.sql.streaming.stateStore.partitions and restarting the query with the same checkpoint location, teams can redistribute saved state across a different number of partitions.
The feature is useful when a workload grows or its resource requirements change. It still requires planning because the query must restart and the saved state must be redistributed before normal processing resumes.
For production use, verify the runtime and RocksDB requirements, measure the existing workload, and monitor the repartition duration. Those steps help you decide whether changing the partition count will solve the actual performance problem.
Join the conversation! Your thoughts help the community grow.