transformWithState を使用した非同期処理（ベータ版）

非同期処理はDatabricks Runtime 19以降、または 環境バージョン6 以上のサーバーレス計算で利用可能です。

Python transformWithStateはasyncioに基づく非同期処理をサポートしています。 状態操作とユーザーロジックを同時にグループキー間で実行し、プロセス間通信をバッチ処理することで、非同期処理はわずかなコード変更のみで同期処理よりも高いスループットを実現します。 このスループット向上はサードパーティの非同期ライブラリを必要としません。 上級者は非同期プログラミングパターンや非同期対応ライブラリでさらにアプリケーションを最適化できます。

非同期処理を使用するには、同期StatefulProcessorの代わりにAsyncStatefulProcessorを実装してください。 AsyncStatefulProcessorAPIは同期StatefulProcessorAPIをミラーリングするため、ほとんどのアプリケーションは非同期APIを利用するためにわずかな変更で済みます。 「AsyncStatefulProcessor の実装」を参照してください。

同期transformWithState APIおよびコアコンセプトについては、「カスタム ステートフル アプリケーションをtransformWithStateで構築する」を参照してください。

AsyncStatefulProcessor の実装

同期 StatefulProcessor を AsyncStatefulProcessorに変換するには、以下の変更を加えてください:

• APIメソッド(init、 close、 handleInputRows、 handleExpiredTimer、 handleInitialState)を async def キーワードで定義します。
• awaitで状態やタイマーの値を読み、更新するか、Pythonのasyncioライブラリを使って実行してください。 これは valueState.get() のような状態操作や registerTimerのようなタイマー操作に当てはまります。 handle.getValueStateのような状態オブジェクトの作成は同期的に行われます。

非同期処理には以下の考慮事項が適用されます。

• アプリケーションがメンバー変数や外部システムにデータを保存している場合、DataBricksは並行実行のためにロジックの書き直しを推奨しています。 handleInputRowsとhandleExpiredTimerはグループキー間で同時に実行できるため、インターリーブされた実行は共有データを破損させてはなりません。 ほとんどのアプリケーションはすでにこの要件を満たしています。
• Databricksは、状態操作によるエラーを捕捉または抑制しないことを推奨しています。 Apache Sparkはこれらのエラーを代わりに処理してくれます。 状態操作が失敗した場合、Apache Spark はタスクを失敗として扱い、再試行します。
  • AsyncStatefulProcessor では、状態操作エラーは自動で管理され、コードに表示されることはありません。
  • 同期 StatefulProcessor では、コード内で状態操作エラーがスローされますが、それを抑制するとデータの正確性が損なわれる可能性があります。

例:各グルーピングキーの行数カウント

以下の例は、各グルーピングキーごとに行数を数える AsyncCountProcessor を定義しています。 value_schema変数はランニングカウントを格納するValueStateのスキーマを定義します。 同期 StatefulProcessor と比べると、変更点は各メソッドの async def キーワードと、状態の読み取りおよび更新操作の await です。 init 内の getValueState の呼び出しは同期のままです。 プロセッサを以下のコードのように定義します:

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

非同期プロセッサでクエリを実行します

非同期プロセッサでクエリを実行するには、AsyncStatefulProcessor を transformWithState に渡します。 クエリは同期パスと同じ構文を使用します。 非同期 API と同期 API は同じ状態フォーマットを共有しているため、既存のクエリを非同期 AsyncStatefulProcessor と同期 StatefulProcessor の間で切り替えても、同じチェックポイントを再利用できます。

例: events サンプルデータセット内のイベントをカウントする

以下の例は events サンプルデータセットに対して AsyncCountProcessor を実行します。 各レコードには time フィールド(エポック秒)と、値が Open または Close である action フィールドがあります。 クエリは action ごとにグループ化し、それぞれのアクションタイプのイベント数をカウントします。 その他のサンプル データセットについては、「サンプル データセット」を参照してください。

input_schema変数はソースレコードのスキーマを定義し、output_schema変数はプロセッサが発行する行のスキーマを定義します。 サンプルデータセットをストリームとして読み取るには、両方のスキーマを定義し、以下のコードのようにクエリを開始します。

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

クエリが完了した後、以下のコードのように各アクションタイプの実行回数を確認します:

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

非同期状態およびタイマー操作

AsyncStatefulProcessorでは、値を読み書きする状態変数やタイマーの操作は非同期です。 これらの操作のほとんどは単一の結果を返し、それを await で取得します。 代わりに、コレクションを返す操作は非同期反復子を返し、それを async for を使用して消費します。 Pythonにおけるasync/awaitおよび非同期イテレーターの入門については、Python asyncioドキュメントを参照してください。

以下の表は、 awaitで取得可能な単一の結果を返す操作を示しています。

以下の表は、非同期反復子を返し、その非同期反復子を async for で取得できる操作を示しています。

例： async for

たとえば、AsyncListState の値を読み取るには、以下のコードのように async for で反復処理します。

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

状態オブジェクトを作成し、状態変数を削除するメソッドは同期的です: getValueState、 getMapState、 getListState、 deleteIfExists。

各状態タイプの説明については「 カスタム状態タイプ」を参照してください。

非同期プログラミング パターンを使用して最適化する

非同期処理は、ロジックがネットワークリクエストなどの外部操作を待つ場合に有用です。 各リクエストを順番に待つ代わりに、 asyncio を使ってリクエストを同時に実行し、アイドル時間を短縮します。

例: 並行してリクエストを実行する asyncio.gather

以下の例では、asyncio.gather を使用してすべての行あたりの HTTP 要求を同時に送信し、完了を待ってから最大スコアを状態に保存します。 プロセッサを以下のコードのように定義します:

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