10.2 Stateless UDFs, Stateful Operators, and StateStores

Decide whether the requirement needs history

A stateless UDF can validate a transaction from that row's fields alone: normalize a code, classify an amount, or check a custom format. Its output does not depend on which transaction happened in a previous batch. A stateful operator instead maintains information across input groups or batches, such as running customer totals, seen event identifiers, window counts or unmatched join records.

The distinction is semantic, not whether the function is written as a Python class. A module-level dictionary inside a scalar UDF is not reliably shared, partitioned or checkpointed. Tasks may run on different executors, Python workers can restart, and failures can replay input. Ordinary UDF instances provide no durable customer-history contract.

Stateless validation example

from pyspark.sql import functions as F
from pyspark.sql.types import BooleanType

@F.udf(returnType=BooleanType())
def valid_reference(value):
    if value is None:
        return False
    return value.startswith("R-") and len(value) == 8

checked = transactions.withColumn("valid_reference", valid_reference("reference"))

The function handles nulls itself and returns a declared BooleanType. It can run on batch or supported streaming DataFrame projections without accumulating earlier references. Prefer built-in string/regex expressions if they express the same rule; the example illustrates a scalar UDF contract. Registering it with spark.udf.register makes it SQL-callable but does not add persistent state.

Built-in stateful operations

A streaming groupBy("customer_id").count() retains a count per key. Streaming dropDuplicates remembers identifiers seen in previous batches. A stream-stream join buffers records while waiting for possible counterparts. These operators use Spark-managed state. Their memory and disk cost depends on key cardinality, retained event history and cleanup policies.

A watermark can let an eligible event-time operator finalize old state, while an application-specific timeout can define expiry for custom state. They are different mechanisms. Processing-time timeout reasons about elapsed processing time; event-time timeout depends on watermark advancement, which may stop if relevant data does not advance. Neither is an ordinary Python sleep or a promise of an exact wall-clock callback deadline.

What a StateStore provides

A StateStore is Spark's versioned keyed storage for stateful streaming tasks. Each micro-batch reads and updates an appropriate state version. Durable checkpoint data enables recovery after failure and helps replay a failed batch consistently. State is partitioned with the operator; it is not one giant Python dictionary on the driver.

The default HDFS-backed provider uses executor memory with checkpoint files in the configured fault-tolerant filesystem. "HDFS-backed" names the provider and does not require the checkpoint path to be an actual HDFS cluster. The RocksDB provider uses native storage and checkpoint mechanisms, often reducing JVM heap pressure for large state. RocksDB still uses memory, local disk and durable checkpoint IO; it does not make state free or eliminate all OOM risks.

Custom keyed state in PySpark

Spark 3.5 provides applyInPandasWithState on grouped streaming DataFrames. The function receives a key, an iterator of pandas batches and a GroupState object. It yields pandas output batches, and the user declares both outputStructType and stateStructType. Do not assume the Scala/Java mapGroupsWithState signature is a direct Python DataFrame method.

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

def running_count(key, batches, state):
    previous = state.get[0] if state.exists else 0
    current = previous + sum(len(pdf) for pdf in batches)
    state.update((current,))
    yield pd.DataFrame({"customer_id": [key[0]], "n": [current]})

counts = events.groupBy("customer_id").applyInPandasWithState(
    running_count,
    outputStructType="customer_id STRING, n LONG",
    stateStructType="n LONG",
    outputMode="Update",
    timeoutConf=GroupStateTimeout.NoTimeout
)

This illustrative example assumes a string customer_id, counts every input row, uses a one-field tuple state and retains it indefinitely. The output writer must support Update. NoTimeout is intentional here: a production high-cardinality stream needs a retention policy. A timeout-aware function must distinguish a timeout invocation from new data and deliberately call state.remove or set new timeout metadata.

Recovery, compatibility and output

State schema and grouping keys are part of checkpoint compatibility. Changing a key type or adding an incompatible state field is not the same as adding a simple projection. Keep separate durable checkpoint locations for separate queries, and plan migrations/replay when changing stateful logic. Deleting a checkpoint loses the history of processed offsets and state.

State recovery does not guarantee an external callback is exactly-once. A retried output can duplicate a database insert unless the sink handles it. Keep state transitions and output semantics deliberate, monitor state row counts and memory/progress metrics, and test null keys, multiple input batches, key expiry and recovery with the supported runtime. See Structured Streaming state and arbitrary stateful operations.