1. Introduction

Stateful Processing in Apache Beam allows a DoFn to read and write persistent state as it processes individual elements. State is scoped per key and per window.

2. Why This Concept Exists

By default, Apache Beam transforms are stateless—workers process each element independently without remembering previous elements. However, advanced business logic requires memory. For example:

Deduplicating duplicate events.
Building session click trails.
Enforcing rate limits (e.g. max 5 API calls per user). Stateful processing provides local, fast, and fault-tolerant storage cells that worker threads can access.

3. Key Terminology

State Cell: An individual, persistent data holder (like a variable) stored on worker disks and synced with backend storage.
ValueState: A state cell that holds a single value (e.g. a counter or last seen timestamp).
BagState: A state cell that holds an append-only collection of values (ideal for buffering records).
StateSpec: The declaration describing the type of state cell (e.g., ValueStateSpec).

4. How It Works

Declare State: Define state cells at the class level of your DoFn using state specifications.
Route by Key: The input PCollection must be key-value pairs (K, V). State is automatically isolated per key.
Read/Write: Inside the process method, read state values using .read() and write updates using .write(new_value) or .add(item).

5. Visual Diagram

User-A process()

Reads/writes state cell dedicated to User-A

User-B process()

Reads/writes state cell dedicated to User-B

6. Code Example

Using ValueState to count elements processed for each specific key:

python
import apache_beam as beam
from apache_beam.state import ValueStateSpec, ReadModifyWriteStateSpec
from apache_beam.coders import VarIntCoder

class StatefulCounterDoFn(beam.DoFn):
    # 1. Declare state cell at class level
    COUNT_STATE = ValueStateSpec("count_state", VarIntCoder())

    def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)):
        # Unpack key-value
        (key, value) = element
        
        # 2. Read current state value (defaults to None if empty)
        current_count = count_state.read() or 0
        
        # 3. Update state
        new_count = current_count + 1
        count_state.write(new_count)
        
        # Yield output
        yield (key, new_count)

7. Code Explanation

COUNT_STATE = ValueStateSpec(...) defines a state cell named "count_state" holding integer coders.
count_state=beam.DoFn.StateParam(COUNT_STATE) injects the state cell parameter into the process method.
count_state.read() retrieves the current key's count.
count_state.write(new_count) updates the count, persisting it for subsequent elements sharing the same key.

8. Real Production Example

In fraud detection, you monitor user logins. You use BagState to store the last 5 login locations for a user key. When a new login occurs, you read the bag state to evaluate geographical jumps, alerting security if a login occurs in New York 10 minutes after a login in Tokyo.

9. Common Mistakes

Applying stateful ParDo to non-keyed collections: Stateful transforms require key-value inputs. Applying it to flat elements throws a validation exception.
Assuming global state scope: State is isolated per key. Workers processing User-A cannot see the state values of User-B.

10. Best Practices

Always specify efficient coders (like VarIntCoder or StrUtf8Coder) to reduce serialization time when syncing state to disk.
Clear state cells using .clear() when they are no longer needed to prevent state bloating and reduce storage costs.

11. Summary

Stateful processing provides persistent, key-scoped variables.
ValueState stores single values; BagState stores append-only lists.
Enforces partition scaling by isolating state by keys.

12. Interactive Challenges

Challenge 1: Stateful Element Counter (Beginner)

Define a DoFn class named KeyCounterDoFn that uses a ValueStateSpec holding integers to count the occurrences of each key, yielding (key, running_count).

Challenge 2: Stateful Deduplicator (Intermediate)

Write a DoFn class named DeduplicateDoFn that uses a ValueStateSpec holding booleans to ensure that only the first occurrence of each key is yielded, discarding subsequent duplicates.

Challenge 3: Transaction Batch Buffer (Advanced)

Define a DoFn class named TransactionBufferDoFn that uses BagStateSpec to collect transaction floats for a store key, and a ValueStateSpec to count them. When the count reaches 3 transactions, it yields the list of values and clears both states.

Stateful Processing