Asynchronous processing with transformWithState (Beta)

Beta

Asynchronous processing for the Python row-based transformWithState API is in Beta. See Databricks preview releases.

Asynchronous processing is available in Databricks Runtime 19 and above.

Python transformWithState supports asynchronous processing built on asyncio. By running state operations and user logic concurrently across grouping keys and batching inter-process communication, asynchronous processing has higher throughput than synchronous processing with only minor code changes. This throughput gain does not require any third-party async libraries. Advanced users can further optimize their applications with async programming patterns and async-enabled libraries.

To use asynchronous processing, implement an AsyncStatefulProcessor instead of the synchronous StatefulProcessor. The AsyncStatefulProcessor API mirrors the synchronous StatefulProcessor API, so most applications require only small changes to use the asynchronous API. See Implement an AsyncStatefulProcessor.

For the synchronous transformWithState API and core concepts, see Build a custom stateful application with transformWithState.

note

Asynchronous processing is available for the Python row-based transformWithState API only. It is not supported for transformWithStateInPandas or for the Scala transformWithState API. Asynchronous processing is not supported in serverless compute.

Implement an AsyncStatefulProcessor

To convert a synchronous StatefulProcessor to an AsyncStatefulProcessor, make the following changes:

• Define the API methods (init, close, handleInputRows, handleExpiredTimer, and handleInitialState) with the async def keyword.
• Read and update state and timer values with await, or run them using Python's asyncio library. This applies to state operations such as valueState.get() and to timer operations such as registerTimer. Creating state objects, such as handle.getValueState, remains synchronous.

The following considerations apply to asynchronous processing:

• If your application stores data in member variables or in external systems, Databricks recommends that you rewrite logic to be safe for concurrent execution. Because handleInputRows and handleExpiredTimer can run concurrently across grouping keys, interleaved runs must not corrupt shared data. Most applications already meet this requirement.
• Databricks recommends that you don't catch or suppress errors from state operations. Apache Spark handles these errors for you. If a state operation fails, Apache Spark fails the task and retries it.
  • In an AsyncStatefulProcessor, state operation errors are managed for you and never surfaced to your code.
  • In a synchronous StatefulProcessor, state operation errors are raised in your code, but suppressing them can compromise data correctness.

Example: count rows for each grouping key

The following example defines an AsyncCountProcessor that counts the number of rows for each grouping key. The value_schema variable defines the schema of the ValueState that stores the running count. Compared to a synchronous StatefulProcessor, the changes are the async def keyword on each method and await on the state read and update operations. The call to getValueState in init remains synchronous. Define the processor as in the following code:

Python

from pyspark.sql import Row
from pyspark.sql.streaming import AsyncStatefulProcessor, AsyncStatefulProcessorHandle
from pyspark.sql.types import StructType, StructField, LongType

value_schema = StructType([StructField("count", LongType(), True)])

class AsyncCountProcessor(AsyncStatefulProcessor):
  async def init(self, handle: AsyncStatefulProcessorHandle) -> None:
    self.count = handle.getValueState("count", value_schema)

  async def handleInputRows(self, key, rows, timerValues):
    total = (await self.count.get() or (0,))[0]
    for _ in rows:
      total += 1
    await self.count.update((total,))
    yield Row(action=key[0], count=total)

  async def close(self) -> None:
    pass

Run a query with an async processor

To run a query with an async processor, pass your AsyncStatefulProcessor to transformWithState. The query uses the same syntax as the synchronous path. The async and synchronous APIs share the same state format, so you can switch an existing query between an AsyncStatefulProcessor and a synchronous StatefulProcessor while reusing the same checkpoint.

Example: count events in the events sample dataset

The following example runs AsyncCountProcessor against the events sample dataset. Each record has a time field (epoch seconds) and an action field with the value Open or Close. The query groups by action and counts the events for each action type. For more sample datasets, see Sample datasets.

The input_schema variable defines the schema of the source records, and the output_schema variable defines the schema of the rows the processor emits. To read the sample dataset as a stream, define both schemas, then start the query as in the following code:

Python

from pyspark.sql.types import StructType, StructField, StringType, LongType

input_schema = StructType([
  StructField("time", LongType(), True),
  StructField("action", StringType(), True),
])

output_schema = StructType([
  StructField("action", StringType(), True),
  StructField("count", LongType(), True),
])

events = (
  spark.readStream.schema(input_schema)
    .option("maxFilesPerTrigger", 10)
    .json("/databricks-datasets/structured-streaming/events")
)

q = (
  events.groupBy("action")
    .transformWithState(
      statefulProcessor=AsyncCountProcessor(),
      outputStructType=output_schema,
      outputMode="Update",
      timeMode="None",
    )
    .writeStream.format("memory")
    .queryName("async_counts")
    .trigger(availableNow=True)
    .start()
)

q.awaitTermination()

After the query completes, view the running count for each action type as in the following code:

Python

display(spark.sql("SELECT action, MAX(count) AS count FROM async_counts GROUP BY action ORDER BY action"))

Async state and timer operations

In an AsyncStatefulProcessor, state variable and timer operations that read or write values are asynchronous. Most of these operations return a single result that you retrieve with await. Operations that return a collection instead return an async iterator that you consume with async for. For an introduction to async/await and asynchronous iterators in Python, see the Python asyncio documentation.

The following table lists operations that return a single result which you can retrieve with await:

Class	Operations that use await
AsyncValueState	exists, get, update, clear
AsyncMapState	exists, getValue, containsKey, updateValue, removeKey, clear
AsyncListState	exists, put, appendValue, appendList, clear
AsyncStatefulProcessorHandle	registerTimer, deleteTimer

Class	Operations that use await
AsyncValueState	exists, get, update, clear
AsyncMapState	exists, getValue, containsKey, updateValue, removeKey, clear
AsyncListState	exists, put, appendValue, appendList, clear
AsyncStatefulProcessorHandle	registerTimer, deleteTimer

The following table lists operations that return an async iterator which you can retrieve with async for:

Class	Operations that use async for
AsyncMapState	iterator, keys, values
AsyncListState	get
AsyncStatefulProcessorHandle	listTimers

Class	Operations that use async for
AsyncMapState	iterator, keys, values
AsyncListState	get
AsyncStatefulProcessorHandle	listTimers

Example: async for

For example, to read the values in an AsyncListState, iterate with async for as in the following code:

Python

total = 0
async for value in self.items.get():
  total += value[0]

The methods that create state objects and delete state variables remain synchronous: getValueState, getMapState, getListState, and deleteIfExists.

For a description of each state type, see Custom state types.

Optimize with async programming patterns

Asynchronous processing is useful when your logic waits on external operations, such as network requests. Instead of waiting for each request in sequence, use asyncio to run the requests concurrently and reduce idle time.

Example: run concurrent requests with asyncio.gather

The following example uses asyncio.gather to fire all per-row HTTP requests concurrently and wait for them to complete, then stores the maximum score in state. Define the processor as in the following code:

Python

import asyncio
import aiohttp
from pyspark.sql import Row
from pyspark.sql.streaming import AsyncStatefulProcessor

class HttpScoreRowGatherProcessor(AsyncStatefulProcessor):
  async def init(self, handle):
    self._score_state = handle.getValueState("last_score", "score double")
    self._session = aiohttp.ClientSession()

  async def _fetch_score(self, row) -> float:
    async with self._session.get(
      f"https://api.example.com/score/{row.event_id}"
    ) as resp:
      return (await resp.json())["score"]

  async def handleInputRows(self, key, rows, timerValues):
    user_id = key[0]
    scores = await asyncio.gather(*[self._fetch_score(row) for row in rows])

    max_score = max(scores)
    await self._score_state.update((max_score,))
    yield Row(user_id=user_id, score=max_score)

  async def close(self):
    await self._session.close()