transformWithState を使用した非同期処理（ベータ版）
ベータ版
Python 行ベース transformWithState API の非同期処理はベータ版です。Databricks プレビューリリースをご覧ください。
非同期処理は、Databricks Runtime 19 以降、または環境バージョン 6 以降のServerless コンピュートで利用できます。
Python transformWithState は、asyncio 上に構築された非同期処理をサポートしています。グループ化キー全体で状態操作とユーザーロジックを並行して実行し、プロセス間通信をバッチ処理することで、非同期処理はわずかなコード変更で同期処理よりも高いthroughputを実現します。このthroughputの向上に、サードパーティの非同期ライブラリは必要ありません。上級ユーザーは、非同期プログラミングパターンと非同期対応ライブラリを使用して、アプリケーションをさらに最適化できます。
非同期処理を使用するには、同期型の StatefulProcessor ではなく AsyncStatefulProcessor を実装してください。AsyncStatefulProcessor API は同期型の StatefulProcessor API をミラーリングしているため、ほとんどのアプリケーションでは、非同期 API を使用するために必要な変更はわずかです。AsyncStatefulProcessorの実装を参照してください。
同期 transformWithState API および基本概念については、transformWithState を使用したカスタム ステートフル アプリケーションの構築を参照してください。
非同期処理は、Python 行ベースの transformWithState API でのみ利用できます。transformWithStateInPandas または Scala transformWithState API ではサポートされていません。
～を実装する AsyncStatefulProcessor
同期 StatefulProcessor を AsyncStatefulProcessor に変換するには、以下の変更を行ってください：
async defキーワードを使用して、API メソッド (init、close、handleInputRows、handleExpiredTimer、およびhandleInitialState) を定義します。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 に渡します。このクエリーは、同期パスと同じ構文を使用します。非同期 APIs と同期 APIs は同じ状態フォーマットを共有しているため、同じチェックポイントを再利用しながら、既存のクエリーを 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 で消費する非同期イテレータを返します。async/await と Python の非同期イテレータの概要については、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()