How do I do stateful aggregation with transformWithState?

Summary

• transformWithState is the arbitrary stateful operator in Structured Streaming since Apache Spark 4.0. You group a stream by key, pass a StatefulProcessor, and keep per-key state in ValueState, ListState or MapState variables that Spark checkpoints with the query.
• It needs the RocksDB state store. Set spark.sql.streaming.stateStore.providerClass to org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider; with another provider the query fails with STORE_BACKEND_NOT_SUPPORTED_FOR_TWS.
• Python has two operators: transformWithStateInPandas (Spark 4.0 and later, pandas DataFrames) and the row-based transformWithState (Spark 4.1 and later). Scala and Java call transformWithState on a groupByKey result.
• For ordinary windowed sums and counts, a built-in aggregation with a watermark is simpler. Use transformWithState when you need custom emission rules, timers, several state variables per key, or TTL-based eviction.
• TTL is set per state variable (ttlDurationMs in Python, TTLConfig in Scala) and works only in processing-time mode.

How do I do stateful aggregation with transformWithState?

Group the stream by key, write a StatefulProcessor whose handleInputRows reads the key's current state, folds in the new rows and writes the result back, then pass it to transformWithState with an output schema, an output mode and a time mode. Spark keeps that state in the RocksDB state store and checkpoints it with the query, so a running total survives restarts. The API names below follow the Apache Spark 4.2.0 Structured Streaming guide; the operator itself has been available since Spark 4.0.

What is transformWithState?

The Spark docs describe it as the next-generation replacement for mapGroupsWithState and flatMapGroupsWithState in Scala and applyInPandasWithState in Python. It's object-oriented: you implement a processor class, and Spark calls handleInputRows once for each grouping key that has new rows, passing every row for that key. The API reference notes that the operator requires a full shuffle.

Building block	What it does	Rule worth knowing
init(handle)	Declares state variables through the StatefulProcessorHandle	You can't register timers here
handleInputRows(key, rows, timerValues)	Processes one key's new rows and emits zero or more rows	Iterate the rows; don't assume a count
handleExpiredTimer	Runs when a timer you registered fires	It can't read input rows
ValueState, ListState, MapState	Per-key state; lists are optimized for appends, maps for point lookups	Each needs a unique name and a schema or encoder
Time mode	None, ProcessingTime or EventTime	Timers and TTL depend on it
TTL	Expires a state variable; any update resets the expiry	Processing-time mode only

Build a running total per key

This processor keeps one running total per user and drops users idle for a day. It uses the row-based Python API (Spark 4.1 and later); on Spark 4.0, use transformWithStateInPandas, which passes pandas DataFrames instead of rows.

from typing import Iterator
from pyspark.sql import Row
from pyspark.sql.streaming import StatefulProcessor, StatefulProcessorHandle

spark.conf.set("spark.sql.streaming.stateStore.providerClass",
               "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider")
DAY_MS = 24 * 60 * 60 * 1000

class RunningTotal(StatefulProcessor):
    def init(self, handle: StatefulProcessorHandle) -> None:
        self.total = handle.getValueState("total", "total BIGINT", ttlDurationMs=DAY_MS)

    def handleInputRows(self, key, rows, timerValues) -> Iterator[Row]:
        current = self.total.get()[0] if self.total.exists() else 0
        current += sum(r.amount for r in rows)
        self.total.update((current,))
        yield Row(user_id=key[0], total=current)

    def close(self) -> None:
        pass

Wire it into the query. TTL requires ProcessingTime mode, and the sink must accept the output mode you pick:

totals = (events.groupBy("user_id")
    .transformWithState(
        statefulProcessor=RunningTotal(),
        outputStructType="user_id STRING, total BIGINT",
        outputMode="Update",
        timeMode="ProcessingTime",
    ))

query = (totals.writeStream.outputMode("update").format("console")
    .option("checkpointLocation", "/tmp/checkpoints/running_total")
    .start())

File sinks accept only Append, while the Console, Kafka, Foreach and ForeachBatch sinks accept Update. In Scala the processor extends StatefulProcessor[K, I, O], creates state in init(outputMode, timeMode) with getHandle.getValueState[Long]("total", Encoders.scalaLong, TTLConfig(java.time.Duration.ofDays(1))), and runs through groupByKey(...).transformWithState(processor, TimeMode.ProcessingTime(), OutputMode.Update()).

Built-in aggregation or transformWithState?

Start with the built-in operators. A groupBy with withWatermark and window handles windowed sums and counts, and Spark cleans up its state for you. The Spark guide positions arbitrary stateful operators for cases that need more than aggregations, such as sessionization with custom state. Choose transformWithState when you need to decide yourself when to emit, keep several state variables per key, fire timers, or expire keys with TTL.

How Databricks approaches transformWithState

On Databricks, transformWithState and transformWithStateInPandas require Databricks Runtime 16.2 or above. RocksDB is the default state store provider from Databricks Runtime 17.3; on 17.2 and below, set the provider as shown above. Real-time mode needs Databricks Runtime 17.3 LTS or above and the row-based transformWithState, because transformWithStateInPandas isn't supported there. Databricks also recommends built-in operators for plain aggregations, deduplication and streaming joins. See Build a custom stateful application.

FAQs

Does transformWithState work without RocksDB?

No. In open-source Spark the operator only supports RocksDBStateStoreProvider, and any other backend fails with STORE_BACKEND_NOT_SUPPORTED_FOR_TWS.

How do I emit a result when a key goes quiet?

Register a timer with handle.registerTimer(...) inside handleInputRows and emit from handleExpiredTimer. Timers are checkpointed with the query and can fire after a restart.

Can I change the state schema later?

Within limits. You can add or remove state variables between runs (call deleteIfExists in init to purge a removed one). Evolving a variable's fields needs spark.sql.streaming.stateStore.encodingFormat set to avro, and renames, type narrowing and key-schema changes aren't supported.

How do I unit-test a processor?

The Spark 4.2.0 guide documents TwsTester, which feeds rows to a processor without running a stream. It doesn't model TTL or automatic watermark advancement, so keep an end-to-end test as well.