transformWithState(ベータ版)による非同期処理

非同期処理はDatabricks Runtime 19以降で利用可能です。

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

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

同期transformWithStateAPIおよびコアコンセプトについては、「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 フィールド(エポック秒)と action フィールドがあり、 Open または Closeの値があります。 クエリはそれぞれのアクションタイプごとにグループ化 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()