Why am I getting OOM in a stateful streaming query?

Summary

• Stateful operators keep state across micro-batches. With the default HDFS-backed state store, that state sits in the JVM memory of the executors, so state that never gets cleaned up ends in long GC pauses and out-of-memory errors.
• Bound the state first: watermarks on aggregations, deduplication and stream-stream joins, and timeouts or TTL in arbitrary stateful code.
• If the state is legitimately large, use the RocksDB state store (Spark 3.2 and above), which keeps state in native memory and on local disk, and cap its memory with boundedMemoryUsage.
• Confirm the cause in the Structured Streaming tab and in stateOperators. RocksDB is the default state store on Databricks Runtime 17.3 and above.

Why am I getting OOM in a stateful streaming query?
Stateful operators keep state across micro-batches, and with the default HDFS-backed provider "the state data is maintained in the JVM memory of the executors". State that only grows causes long garbage-collection pauses and, in the Databricks docs' words, "out-of-memory (OOM) errors because of the amount of data kept in state". Bound the state first, then move it off the heap with RocksDB, then size the executors.

Which operators keep state forever?

Operator	State grows without bound when	What bounds it, per the docs
Aggregation	The watermark does not apply	withWatermark on the event-time column the aggregation uses, called before it, in Append or Update mode
dropDuplicates	There is no watermark: "the query stores the data from all the past records as state"	A watermark plus dropDuplicates(["guid", "eventTime"]), or dropDuplicatesWithinWatermark
Stream-stream inner join	"all past input must be saved"	Watermarks on both inputs plus an event-time constraint, such as a time range join condition
flatMapGroupsWithState, mapGroupsWithState, applyInPandasWithState	No timeout is set and remove() is never called	A timeout set on every call, and remove() when it fires

Why is my watermark not clearing state?

The Spark docs give four conditions for watermarking to clean aggregation state. The output mode must be Append or Update, because Complete mode "requires all aggregate data to be preserved". "The aggregation must have either the event-time column, or a window on the event-time column." withWatermark must be called on that same column. It must also be called before the aggregation. A query that sets a watermark and then groups only by a business key does not meet the second condition, so its state is never cleaned. A larger watermark delay keeps more state: Databricks notes that "because of the larger state size, a larger threshold might also require additional computing resources". The docs' working example groups by a window:

from pyspark.sql.functions import window

windowedCounts = words \
    .withWatermark("timestamp", "10 minutes") \
    .groupBy(window(words.timestamp, "10 minutes", "5 minutes"), words.word) \
    .count()

How do I expire state in arbitrary stateful code?

In Python the operator is applyInPandasWithState. Set a timeout on every call and remove state when it fires:

import pandas as pd
from pyspark.sql.streaming.state import GroupStateTimeout

def count_events(key, pdf_iter, state):
    if state.hasTimedOut:
        state.remove()
        return
    total = state.get[0] if state.exists else 0
    for pdf in pdf_iter:
        total += len(pdf)
    state.update((total,))
    state.setTimeoutDuration(30 * 60 * 1000)
    yield pd.DataFrame({"id": [key[0]], "events": [total]})

# events is a streaming DataFrame with an id column
counts = events.groupBy("id").applyInPandasWithState(
    count_events,
    outputStructType="id long, events long",
    stateStructType="total long",
    outputMode="Update",
    timeoutConf=GroupStateTimeout.ProcessingTimeTimeout,
)

The Scala GroupState docs say the timeout "is reset every time the function is called on a group", so "the user has to set the timeout duration every time the function is called, otherwise, there will not be any timeout set". A processing-time timeout also needs a trigger to fire: with no data in the stream, "there will not be any trigger and timeout function call will not occur until there is data". Since Spark 4.0 the docs encourage transformWithState, which supports "automatic TTL based eviction".

How do I move state off the JVM heap?

For legitimately large state, the docs describe RocksDB (Spark 3.2 and above): "Rather than keeping the state in the JVM memory, this solution uses RocksDB to efficiently manage the state in the native memory and the local disk."

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

If left unbounded, RocksDB memory use "could grow indefinitely and potentially cause OOM (out-of-memory) issues". Set spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage to true (default false) and spark.sql.streaming.stateStore.rocksdb.maxMemoryUsageMB (default 500); the docs call this a soft limit. The provider class cannot change for an existing query: "To read the previous state of the query properly, the class of state store provider should be unchanged." Switching means discarding the checkpoint and starting a new query. The same applies to spark.sql.shuffle.partitions.

How do I confirm state is the problem?

Open the Structured Streaming tab in the Spark web UI, not the Streaming (DStreams) tab. It reports "Aggregated Number Of Total State Rows", "Aggregated State Memory Used In Bytes" and "Aggregated Number Of State Rows Dropped By Watermark". In code, query.lastProgress includes stateOperators, with numRowsTotal, numRowsRemoved, numRowsDroppedByWatermark and memoryUsedBytes. A total that keeps rising points to state that is not being cleaned. With the default provider the state lives in the executor heap, so if it is bounded but still too large, raise spark.executor.memory (default 1g), for example with --conf spark.executor.memory=8g.

How does Databricks handle this?

RocksDB is the default state store provider in Databricks Runtime 17.3 and above. On earlier versions, set spark.sql.streaming.stateStore.providerClass to com.databricks.sql.streaming.state.RocksDBStateStoreProvider before starting the query. Databricks recommends changelog checkpointing for all Structured Streaming stateful queries, and it is on by default from Databricks Runtime 17.3. In Lakeflow Pipelines, use watermarks for "aggregations, joins, and deduplication".

FAQs

Does checkpointing reduce the memory state uses?

No. Checkpointing saves state for recovery, while the default provider keeps the live state in a memory map "and then backed by files". RocksDB changes where the live state is kept.

Does Lakeflow Pipelines have the same problem?

It can. Databricks documents watermarks for stateful processing in pipelines to avoid problems such as out-of-memory errors.