Streaming state management: patterns, pitfalls and fixes

Streaming state management: patterns, pitfalls and fixes

By day 99 of a serious data engineering journey, you've likely moved beyond simple stateless transformations. Stateful streaming queries—where you maintain windows, aggregations, or join state across batches—are where complexity lives and where production systems fail silently. We'll walk through the patterns that keep your state manageable, the watermarks that bound it, the storage engines that scale it, and the monitoring that saves you at 3 AM.

Why state matters in streaming

Stateless operations—filtering, mapping, exploding rows—are fire-and-forget. You process a record, emit a result, move on. State is different. When you aggregate sales by region over a 10-minute window, or join a stream of clicks to a stream of impressions, Spark must remember intermediate results. That state lives in memory or on disk. It grows. It can leak. It can cause your driver or executors to run out of heap and crash the entire job.

Most teams learn this the hard way: a query that ran fine for weeks suddenly grinds to a halt or OOMs. The culprit is usually unbounded state—results accumulating with no mechanism to clean them up.

Aggregations and joins keep state

Two main patterns generate state in Spark Structured Streaming: aggregations and stream-to-stream joins.

Windowed aggregations

A windowed aggregation groups events by time and dimension, then computes a metric (count, sum, average) within each window. Spark must store all events that fall into overlapping windows until the window closes and results are emitted.

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, TimestampType

# Schema for incoming events
schema = StructType([
    StructField("event_id", StringType()),
    StructField("user_id", StringType()),
    StructField("region", StringType()),
    StructField("revenue", IntegerType()),
    StructField("timestamp", TimestampType())
])

# Read from Kafka or socket
events = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "events") \
    .load() \
    .select(F.from_json(F.col("value").cast("string"), schema).alias("data")) \
    .select("data.*")

# 10-minute windows, 5-minute slide
query = events \
    .groupBy(
        F.window(F.col("timestamp"), "10 minutes", "5 minutes"),
        F.col("region")
    ) \
    .agg(
        F.sum("revenue").alias("total_revenue"),
        F.count("*").alias("event_count")
    ) \
    .writeStream \
    .format("console") \
    .option("checkpointLocation", "/tmp/checkpoint") \
    .start()

query.awaitTermination()

In this example, each sliding window holds state until it closes. A 10-minute window with 5-minute slides means state for the same region overlaps across two windows. Spark tracks which records belong to which windows, aggregates them, and only discards state when the window is definitively past the watermark.

Stream-to-stream joins

Joining two streams is trickier. Spark cannot assume one stream is complete before the other arrives. It must buffer both sides, potentially indefinitely. A click might arrive seconds before its matching impression, or vice versa. Without explicit handling, Spark holds both sides in state forever.

# Two streams: clicks and impressions
clicks = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "clicks") \
    .load() \
    .select(F.from_json(F.col("value").cast("string"), schema).alias("data")) \
    .select("data.*")

impressions = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "impressions") \
    .load() \
    .select(F.from_json(F.col("value").cast("string"), schema).alias("data")) \
    .select("data.*")

# Join on impression_id; state grows on both sides
result = clicks.join(
    impressions,
    clicks.impression_id == impressions.id,
    "inner"
) \
    .select(
        clicks.user_id,
        clicks.click_time,
        impressions.creative_id
    )

query = result.writeStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("topic", "click_impression_join") \
    .option("checkpointLocation", "/tmp/checkpoint_join") \
    .start()

Both sides of the join maintain state: unmatched clicks waiting for impressions, and unmatched impressions waiting for clicks. Without watermarking, these buffers can grow unboundedly, consuming gigabytes of memory within hours.

Watermarks bound state size

A watermark is a timestamp that signals "no more events earlier than this time will arrive." Spark uses it to garbage-collect state for old windows or old join keys. Setting a good watermark is crucial.

How watermarks work

When you define a watermark, you're telling Spark: "Events might arrive late by up to X minutes, but nothing will arrive later than that." Spark tracks the maximum timestamp it has seen and lags the watermark behind it by your specified duration. Once the watermark passes a window's end time, Spark can safely discard that window's state.

# Set a watermark: allow 15 minutes of lateness
events_with_watermark = events \
    .withWatermark("timestamp", "15 minutes")

result = events_with_watermark \
    .groupBy(
        F.window(F.col("timestamp"), "10 minutes", "5 minutes"),
        F.col("region")
    ) \
    .agg(
        F.sum("revenue").alias("total_revenue"),
        F.count("*").alias("event_count")
    )

query = result.writeStream \
    .format("console") \
    .option("checkpointLocation", "/tmp/checkpoint") \
    .outputMode("update") \
    .start()

In this query, the watermark lags 15 minutes behind the latest timestamp. Once a window closes (its end time is earlier than watermark), Spark drops it. If you receive an event with timestamp 15 minutes older than the current watermark, it is silently dropped—a late-arriving record that will not update any window.

Choosing the right watermark delay

A watermark too short (say, 1 minute) will drop valid late data. A watermark too long (say, 24 hours) will bloat state unnecessarily. We often pick based on SLA: if 99% of events arrive within 10 minutes, set a 15-minute watermark and accept the 1% loss, or implement a late-data path that feeds to a correction table.

For stream-to-stream joins, watermarks on both sides are equally important:

clicks_wm = clicks.withWatermark("click_time", "10 minutes")
impressions_wm = impressions.withWatermark("impression_time", "10 minutes")

result = clicks_wm.join(
    impressions_wm,
    clicks.impression_id == impressions.id,
    "inner"
)

query = result.writeStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("topic", "joined") \
    .option("checkpointLocation", "/tmp/checkpoint_join") \
    .start()

Both sides garbage-collect state beyond their watermarks. A click unmatched after 10 minutes past the click timestamp is discarded; same for impressions. Without these watermarks, both sides grow indefinitely.

RocksDB state store for large state

By default, Spark holds state in memory (StateStore backed by ConcurrentHashMap). For queries with large state—millions of keys or gigabytes of aggregates—memory runs out. RocksDB is an embedded key-value store that uses local disk, spilling state to SSD while keeping a warm cache in memory.

Enabling RocksDB

Configure it at the SparkSession level:

spark = SparkSession.builder \
    .appName("streaming_with_rocksdb") \
    .config("spark.sql.streaming.stateStore.providerClass", 
            "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider") \
    .config("spark.sql.streaming.rocksdb.blockCacheSizeMB", 512) \
    .config("spark.sql.streaming.rocksdb.blockSize", 16384) \
    .getOrCreate()

# Now all stateful queries use RocksDB
events_with_watermark = events.withWatermark("timestamp", "15 minutes")

result = events_with_watermark \
    .groupBy(
        F.window(F.col("timestamp"), "10 minutes", "5 minutes"),
        F.col("region")
    ) \
    .agg(
        F.sum("revenue").alias("total_revenue")
    )

query = result.writeStream \
    .format("console") \
    .option("checkpointLocation", "/tmp/checkpoint") \
    .start()

RocksDB configuration tuning

Key settings affect performance:

• blockCacheSizeMB: In-memory cache for frequently accessed blocks. Larger = fewer disk hits, but uses heap. Typical: 256–1024 MB.
• blockSize: Size of each block written to disk. Smaller blocks = more overhead; larger blocks = worse seek performance. Default 4 KB is usually fine; 16 KB for larger state items.
• compactInterval: Frequency of background compaction. Lower values = faster queries but more CPU.

For a join with millions of unmatched keys, RocksDB keeps working keys in cache and spills cold state to disk. This lets you process state sizes that would OOM with in-memory storage.

Monitor state size

State growth is invisible until it kills your job. Monitoring state metrics is non-negotiable in production.

Spark UI and metrics

The Spark UI shows state metrics under the Structured Streaming tab: total state size, number of rows, and whether RocksDB is evicting to disk. Enable verbose logging to catch state explosions early:

# In log4j.properties or at runtime:
# log4j.logger.org.apache.spark.sql.execution.streaming.state=DEBUG

import logging
logger = logging.getLogger("org.apache.spark.sql.execution.streaming.state")
logger.setLevel(logging.DEBUG)

Exporting metrics to monitoring systems

Spark's built-in metrics can be exported to Prometheus, CloudWatch, or other systems. Enable the metric sink and collect state-related gauges:

spark = SparkSession.builder \
    .appName("streaming_with_metrics") \
    .config("spark.metrics.conf", "/path/to/metrics.conf") \
    .getOrCreate()

# metrics.conf includes:
# spark.metrics.namespace=my_app
# driver.source.jvm.class=org.apache.spark.metrics.source.JvmSource
# driver.sink.prometheus.class=org.apache.spark.metrics.sink.PrometheusSink

Write custom code to poll the streaming query for state metrics:

import time

query = result.writeStream \
    .format("console") \
    .option("checkpointLocation", "/tmp/checkpoint") \
    .start()

while query.isActive:
    status = query.status
    progress = query.lastProgress
    
    if progress:
        state_memory = progress.get("stateMemory", 0)
        state_rows = progress.get("numStateRows", 0)
        
        print(f"State memory: {state_memory / 1e6:.2f} MB")
        print(f"State rows: {state_rows}")
        
        # Alert if state memory exceeds threshold
        if state_memory > 5e9:  # 5 GB
            print("WARNING: State memory too high!")
    
    time.sleep(10)

Alerting on state anomalies

Set up alerts for:

• State memory growth rate accelerating (sign of a watermark issue).
• RocksDB eviction rate spiking (queries slower; consider tuning blockCacheSizeMB).
• Number of state rows not declining after a window close (watermark may not be working).

Common mistakes

Forgetting the watermark: Queries run fine for days, then memory explodes on day 8. State accumulates for every unique key ever seen. Always add a watermark to stateful queries. No exceptions.

Setting the watermark too long: A 7-day watermark on a 10-minute window means state for 1,000 unique keys per minute multiplied by 7 days. That's 10 million keys. For low-cardinality dimensions, this works; for high-cardinality (user IDs, session IDs), it balloons state. Understand your cardinality.

Ignoring join order: When joining two high-cardinality streams, the order matters. Spark buffers both sides, so join the lower-cardinality stream first to reduce memory pressure.

Not testing failover: State is checkpointed, but recovery is slow if state is huge. Test that your job recovers within your SLA. If state is 50 GB, recovery might take minutes, causing data loss if another failure occurs.

Tuning RocksDB without baselines: Don't guess. Measure state size under your typical and peak load. Then adjust blockCacheSizeMB and blockSize based on observed behavior, not defaults.

Key takeaways

• Aggregations and joins generate state; unbounded state crashes production jobs.
• Watermarks are the primary defense: they tell Spark when to garbage-collect old state. Set them conservatively; accept a small percentage of late data loss rather than running out of memory.
• RocksDB scales state beyond memory by using local SSD. Configure blockCacheSizeMB based on your heap and load pattern.
• Monitor state size, growth rate, and RocksDB eviction. Set alerts for anomalies. State explosions are predictable if you watch the metrics.
• Test failover with realistic state sizes. A job that recovers in 30 seconds with 1 GB of state might take 10 minutes with 100 GB.