Traitement asynchrone avec transformWithState

Le traitement asynchrone pour l’API Python basée transformWithState sur les lignes nécessite Databricks Runtime 19 et supérieur, ou le calcul serverless avec l’environnement version 6 et supérieure.

Pour utiliser un traitement asynchrone dans un pipeline Lakeflow, configurez le pipeline pour utiliser le canal PREVIEW. Consultez channel les configurations de pipeline.

Python transformWithState prend en charge le traitement asynchrone construit sur asyncio. En exécutant simultanément les opérations d’état et la logique utilisateur sur l’ensemble des clés de groupement et en regroupant la communication inter-processus, le traitement asynchrone a un débit plus élevé que le traitement synchrone, avec seulement des modifications mineures de code. Ce gain de débit ne nécessite aucune bibliothèque asynchrone tierce. Les utilisateurs avancés peuvent optimiser davantage leurs applications grâce à des schémas de programmation asynchrones et des bibliothèques compatibles async.

Pour utiliser un traitement asynchrone, implémentez un AsyncStatefulProcessor au lieu du synchrone StatefulProcessor. L’API AsyncStatefulProcessor reflète l’API synchrone StatefulProcessor , donc la plupart des applications nécessitent seulement de petits changements pour utiliser l’API asynchrone. Voir « Implémenter un AsyncStatefulProcessor ».

Pour l’API synchrone transformWithState et les concepts de base, voir Construire une application avec état personnalisée avec transformWithState.

Implémentez un AsyncStatefulProcessor

Pour convertir un processus synchrone StatefulProcessor en un AsyncStatefulProcessor, effectuez les modifications suivantes :

• Définissez les méthodes API (init, close, handleInputRows, handleExpiredTimer, et handleInitialState) avec le async def mot-clé.
• Lisez et mettez à jour les valeurs d’état et de minuterie avec await, ou exécutez-les à l’aide de la bibliothèque Python asyncio. Cela s’applique aux opérations d’état telles que valueState.get() et aux opérations du minuteur telles que registerTimer. La création d’objets d’état, tels que handle.getValueState, reste synchrone.

Les considérations suivantes s’appliquent au traitement asynchrone :

• Si votre application stocke les données dans des variables membres ou dans des systèmes externes, Databricks recommande de réécrire la logique pour qu’elle soit sûre pour une exécution concurrente. Parce que handleInputRows et handleExpiredTimer peuvent s’exécuter simultanément à travers les clés de groupement, les exécutions entrelacées ne doivent pas corrompre les données partagées. La plupart des applications remplissent déjà cette exigence.
• Databricks recommande de ne pas intercepter ni masquer les erreurs issues des opérations d’état. Apache Spark gère ces erreurs pour vous. Si une opération d’état échoue, Apache Spark fait échouer la tâche et la réessaie.
  • Dans un AsyncStatefulProcessor, les erreurs d’opérations d’état sont gérées pour vous et ne sont jamais exposées à votre code.
  • Dans une opération d’état synchrone StatefulProcessor, des erreurs d’opération d’état sont générées dans votre code, mais les supprimer peut compromettre la correction des données.

Exemple : compter les lignes pour chaque clé de regroupement

L’exemple suivant définit un AsyncCountProcessor qui compte le nombre de lignes pour chaque clé de groupement. La variable value_schema définit le schéma du ValueState qui stocke le décompte courant. Comparé à un(e) StatefulProcessor synchrone, les changements sont le mot-clé async def sur chaque méthode et await sur les opérations de lecture et de mise à jour de l’état. L’appel à getValueState dans init reste synchrone. Définissons le processeur comme dans le code suivant :

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

Lancer une requête avec un processeur asynchrone

Pour lancer une requête avec un processeur asynchrone, passez votre AsyncStatefulProcessor à transformWithState. La requête utilise la même syntaxe que le chemin synchrone. Les API asynchrones et synchrones partagent le même format d’état, vous pouvez donc basculer une requête existante entre un AsyncStatefulProcessor et un StatefulProcessor synchrone tout en réutilisant le même point de contrôle.

Exemple : compter les événements dans le jeu de données d’exemple events

L’exemple suivant exécute AsyncCountProcessor sur l’ensemble de données d’exemple events. Chaque enregistrement possède un champ time (secondes d’époque) et un champ action avec la valeur Open ou Close. La requête regroupe par action et compte les événements pour chaque type d’action. Pour plus d’exemples de jeux de données, voir Exemples de jeux de données.

La input_schema variable définit le schéma des enregistrements sources, et la output_schema variable définit le schéma des lignes émises par le processeur. Pour lire l’échantillon de jeu de données comme un flux, définissez les deux schémas, puis lancez la requête comme dans le code suivant :

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

Une fois la requête terminée, consultez le décompte en cours pour chaque type d’action comme indiqué dans le code suivant :

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

Opérations d’état asynchrone et de minuterie

Dans un AsyncStatefulProcessor, les opérations de variable d’état et de minuterie qui lisent ou écrivent des valeurs sont asynchrones. La plupart de ces opérations retournent un seul résultat que vous récupérez avec await. Les opérations qui retournent une collection renvoient plutôt un itérateur asynchrone que vous consommez avec async for. Pour une introduction aux async/await et itérateurs asynchrones en Python, voir la documentation Python asyncio.

Le tableau suivant liste les opérations qui retournent un seul résultat que vous pouvez récupérer avec await:

Le tableau suivant liste les opérations qui retournent un itérateur asynchrone que vous pouvez récupérer avec async for:

Exemple : async for

Par exemple, pour lire les valeurs dans un AsyncListState, itère avec async for comme dans le code suivant :

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

Les méthodes qui créent les objets d’état et suppriment les variables d’état restent synchrones : getValueState, getMapState, getListState, et deleteIfExists.

Pour une description de chaque type d’état, voir Types d’état personnalisés.

Optimiser avec des schémas de programmation asynchrones

Le traitement asynchrone est utile lorsque votre logique attend des opérations externes, telles que les requêtes réseau. Au lieu d’attendre chaque requête en séquence, utilisez asyncio pour exécuter les requêtes simultanément et réduire le temps d’inactivité.

Exemple : exécuter des requêtes concurrentes avec asyncio.gather

L’exemple suivant utilise asyncio.gather pour lancer simultanément toutes les requêtes HTTP par ligne et attendre qu’elles se terminent, puis stocke le score maximal dans l’état. Définissons le processeur comme dans le code suivant :

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