Stateful operators and StateStores with transformWithState

1.Stateful queries and the state store

A stateless streaming query only tracks which source rows it has already processed. A stateful query also keeps intermediate state information and updates it incrementally with each micro-batch. That covers streaming aggregations, distinct, dropDuplicates, stream-stream joins, and custom stateful applications. The state lives in a state store. Its provider is set by spark.sql.streaming.stateStore.providerClass. RocksDB is the default provider in Databricks Runtime 17.3 and above, and on earlier runtimes you have to configure it yourself. Databricks also recommends RocksDB with changelog checkpointing for stateful streams.

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

The checkpoint fixes the shuffle partition count in the same way. Changing spark.sql.shuffle.partitions has no effect on a query that already has a checkpoint. Before you write custom state logic at all, note what Databricks advises: for aggregations, deduplication, and streaming joins, use the built-in operators. Custom stateful processors are for logic those operators can't express.

Sources12

2.Defining a StatefulProcessor and its state variables

With transformWithState, you write a class that extends StatefulProcessor. Spark passes a StatefulProcessorHandle to your init method, and you use that handle to create state variables with getValueState, getListState, or getMapState. Each state variable needs a unique name and a schema. In Python the schema is mandatory, and in Scala you can pass an Encoder instead. You can optionally add a TTL in milliseconds. A MapState needs separate schemas for its keys and its values. One processor can hold several state variables. The grouping key comes from the groupBy before the operator, and Spark keeps every state variable separately for each grouping key.

class SessionTracker(StatefulProcessor):
  def init(self, handle: StatefulProcessorHandle) -> None:
    self.sessions = handle.getMapState("sessions", "session_id string", "count long")

  def handleInputRows(self, key, rows: Iterator[Row], timerValues) -> Iterator[Row]:
    for row in rows:
      session_key = (row["session_id"],)  # session_id is the MapState key
      count = self.sessions.getValue(session_key)[0] if self.sessions.containsKey(session_key) else 0
      new_count = count + 1
      self.sessions.updateValue(session_key, (new_count,))
    yield from []

  def close(self) -> None:
    pass

df.groupBy("user_id").transformWithState(SessionTracker(), ...) # user_id is the grouping key

Sources2

3.handleInputRows, handleExpiredTimer and handleInitialState

Your processing logic goes in its methods. handleInputRows runs for each grouping key that has rows in the micro-batch. It gets those rows as an iterator, can read and write state, and yields output rows. handleExpiredTimer runs when a timer fires, whether or not the key received new rows. Timers let you go beyond simple state eviction, including emitting rows. handleInitialState is optional and pre-populates state before any input rows arrive.

Sources2

4.Invoking the processor on a streaming DataFrame

PySpark offers two operators. transformWithState is row-based, and transformWithStateInPandas passes pandas DataFrames to your handlers. Scala supports only the row-based API. In Python you call either operator on df.groupBy(key) and supply the processor instance, an output schema, an output mode, and a time mode. The SCD type 1 example below keeps only the latest location per user in a ValueState.

        # Initialize the state to store the latest location for each user
        self.latest_location = handle. ? ("latestLocation", value_state_schema)

q = (
    df.groupBy("user")
    .transformWithStateInPandas(
        statefulProcessor=SCDType1StatefulProcessor(),
        outputStructType=output_schema,
        outputMode="Update",
        timeMode="None",
    )
    .writeStream.format("memory")
    .queryName("scd1_output")
    .option("checkpointLocation", f"/tmp/checkpoint_{uuid.uuid4()}")
    .trigger(availableNow=True)
    .start()
)

Sources2

5.Legacy arbitrary stateful operators

Before transformWithState, custom state ran through mapGroupsWithState and flatMapGroupsWithState, plus applyInPandasWithState in Python. In these operators the state update function receives the previous state as a GroupState object. Both Scala operators can accept a user-defined initial state, so a stream started without a valid checkpoint doesn't have to reprocess data. Databricks now recommends transformWithState for arbitrary state transformations. It is also the custom operator that works in pipelines that chain several stateful operators, in DBR 16.2 and later. Those pipelines don't support the legacy operators and allow only append output mode.

def flatMapGroupsWithState[S: Encoder, U: Encoder](
    outputMode: OutputMode,
    timeoutConf: GroupStateTimeout,
    initialState: KeyValueGroupedDataset[K, S])(
    func: (K, Iterator[V], GroupState[S]) => Iterator[U])

Sources213

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.Appending items to a ListState refreshes the TTL of the list.Why is that wrong?

   Each list value has its own TTL, and only a put operation, which overwrites the whole list, resets it.

   Covered in Defining a StatefulProcessor and its state variables

2. 2.You can update a single field of a struct held in ValueState.Why is that wrong?

   A ValueState holds a single value per key, and your logic has to replace the entire value.

   Covered in Defining a StatefulProcessor and its state variables

3. 3.Raising spark.sql.shuffle.partitions and restarting a checkpointed stateful query re-partitions its state.Why is that wrong?

   The partition count is fixed when the checkpoint is created. A query that already has a checkpoint keeps its original count.

   Covered in Stateful queries and the state store

Practise it for real

Run the Databricks SCD type 1 example and confirm that the state store keeps only the latest location per user.

1. 1.Set spark.sql.streaming.stateStore.providerClass to the RocksDBStateStoreProvider class (required below DBR 17.3).

   Why: transformWithState needs RocksDB as the state store provider on older runtimes.

   You should see: The config call returns without error.

2. 2.Define SCDType1StatefulProcessor, with init calling handle.getValueState("latestLocation", value_state_schema).

   Why: A ValueState holds exactly one row per user, which is what SCD type 1 needs.

   You should see: The class definition compiles in the notebook.

3. 3.Create main.stateful_examples, seed scd1_source with rows (u1,1,NYC), (u1,3,SF), (u1,2,LA), (u2,5,London), and read it with spark.readStream.table.

   Why: This gives you a small streaming source with out-of-order times for u1.

   You should see: df is a streaming DataFrame.

4. 4.Run df.groupBy("user").transformWithStateInPandas(...) writing to the memory sink scd1_output with trigger(availableNow=True), then call q.awaitTermination().

   Why: This invokes the processor once per grouping key and saves state under the checkpoint location.

   You should see: The query processes the available data and stops.

5. 5.Query SELECT user, time, location FROM scd1_output ORDER BY user.

   Why: The output shows what the processor emitted from state.

   You should see: u1 -> SF (time 3) and u2 -> London (time 5).