Processamento assíncrono com transformWithState chega à GA no Databricks

O processamento assíncrono para a API transformWithState baseada em linhas em Python agora está em disponibilidade geral (GA). O recurso executa operações de estado e lógica de processamento de forma concorrente via asyncio, reduzindo o tempo ocioso e aumentando a vazão do streaming.

Para que serve

Imagine que você tem um pipeline de streaming com estado (stateful) que consome eventos por chave de agrupamento, atualiza métricas e precisa consultar APIs externas ou gravar estados intermediários com frequência. No modo síncrono, cada operação de estado bloqueia a execução até ser finalizada.

Com o modo assíncrono, você consegue:

• Processar chaves de agrupamento de forma concorrente sem depender de bibliotecas assíncronas externas.
• Fazer chamadas concorrentes para APIs externas usando asyncio.gather dentro da sua lógica de processamento.
• Agrupar em lote a comunicação entre processos, elevando a taxa de transferência (throughput).

Como usar

Para migrar do modelo síncrono, substitua a classe base StatefulProcessor por AsyncStatefulProcessor, declare os métodos com async def e aplique await nas leituras e gravações de estado.

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

O que acontece nesse código:

1. A criação da variável de estado via getValueState dentro do init permanece síncrona.
2. O método handleInputRows passa a ser uma corrotina assíncrona (async def).
3. As operações de leitura (self.count.get()) e persistência (self.count.update()) utilizam a palavra-chave await.

Para executar a query de streaming, chame transformWithState normalmente passando a instância assíncrona:

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()

O que acontece nesse código:

1. A API de chamada transformWithState aceita diretamente o AsyncCountProcessor sem alterar os parâmetros de execução.
2. Como o formato de checkpoint é idêntico entre o processamento síncrono e assíncrono, é possível reutilizar o mesmo diretório de checkpoint de jobs já existentes.

Limitações e requisitos

• Requer Databricks Runtime 19 ou superior, ou computação serverless com versão de ambiente 6 ou superior.
• Em pipelines do Lakeflow, exige a configuração do pipeline no canal PREVIEW.
• Suporte exclusivo para a API em Python baseada em linhas (transformWithState). Não há suporte para transformWithStateInPandas nem para a API em Scala.
• Operações que retornam coleções de dados (como AsyncListState.get ou AsyncMapState.keys) retornam iteradores assíncronos e exigem iteração com async for.