PySparkにおけるArrow UDFの導入: Pandas UDFに代わる、より高速で効率的な選択肢

はじめに

Pythonユーザー定義関数（UDF）は不可欠な拡張メカニズムですが、従来、行ベースの実行による高いオーバーヘッドに悩まされていました。Apache Spark™では、Pandas UDFはArrowベースのシリアル化とバッチ処理を導入することでこの問題の一部を解決し、スカラーPython UDFと比較してスループットを大幅に向上させました。

しかし、Pandas UDFには依然として根本的な制限があります。

• Pandas/Arrowデータ変換は追加のデータコピーを発生させます。ゼロコピーのアプローチは、特定の限られたケースでのみ可能です。例えば、NULL値を持つ列はディープコピーを引き起こします。
• 複雑なデータ型は十分にサポートされていません。例えば、ネストされたStructTypeインスタンスは、集計ユースケースの出力型ではサポートされていません。

Pandas/Arrowデータ変換をなくすことで、Arrow UDFはPandas UDFよりも高速に実行され、メモリ消費量が少なく、より優れたデータ型サポートを提供します。

ネイティブArrow UDF

Databricks Runtime 18.0（リリースノート）から、パフォーマンスの高いUDF実行における画期的な進歩であるネイティブArrow UDFを導入できることを大変嬉しく思います。

ネイティブArrow UDFは、入力をPandasまたはNumPyオブジェクトに変換することなく、Arrowデータ上で直接動作します。これにより、カラムナーレイアウトがエンドツーエンドで保持され、不要なデータコピーが回避され、UDFはArrowのネイティブな計算およびメモリモデルを活用してベクトル化された処理を使用できます。

Arrow UDFを定義するには、ユーザーは新しいPythonデコレータ@arrow_udfを使用でき、戻り値の型とオプションの評価型を指定します。例えば、次のようになります。

ユーザーは、既存のデコレータ@udfと完全な型ヒントを使用して定義することもできます。例えば、次のようになります。

注：関数定義には、すべての引数と戻り値の型ヒントを含める必要があります。
この設計は、スカラーPython UDFのインターフェースと一致しており、スカラーPython UDFにすでに慣れているユーザーに一貫性のある直感的なエクスペリエンスを提供します。

以下にArrow UDFの使用方法を示します。

Pythonでの使用法：

SQLでの使用法：

Arrow UDFインターフェースのバリアントをサポートしています。これには、スカラー関数、集計関数、テーブル関数が含まれます。データフレームAPIでは、Arrow UDFを使用するためのmapInArrowとapplyInArrowも提供しています。次に、これらを一つずつ紹介します。

Arrowスカラー関数

Arrowスカラー関数は行ごとの変換を実行します。これらはスカラーPandas UDFのArrow版であり、df.select()やdf.withColumn()のように列式が期待されるあらゆる場所で使用できます。3つの入力モードがサポートされています：直接、イテレータ、および複数の配列のイテレータです。イテレータのバリアントは、UDFが高価な1回限りの初期化（例：モデルのロードや正規表現パターンのコンパイル）を必要とする場合に役立ちます。セットアップコストがすべてのバッチに償却されるためです。すべての場合において、出力行数は入力行数と一致する必要があります。

• 配列から配列へ：1つ以上のpyarrow.Arrayを受け取り、1つのpyarrow.Arrayを返します。入力配列と出力配列は同じ数の値を持つ必要があります。

• 配列のイテレータから配列のイテレータへ：pyarrow.Arrayのイテレータを受け取り、pyarrow.Arrayのイテレータを返します。このタイプは、UDFの実行に高価な初期化が必要な場合に役立ちます。

• 複数の配列のイテレータから配列のイテレータへ：複数のpyarrow.Arrayのタプルのイテレータを受け取り、pyarrow.Arrayのイテレータを返します。

Arrow集計関数

Arrow集計関数は、1つ以上のpyarrow.Array入力を受け取り、スカラー値を返します。これにより、行のグループを単一の結果に集約します。これらはグループ化された集計Pandas UDFのArrow版であり、groupBy().agg()またはウィンドウ操作で使用されます。スカラー関数と同様に、集計関数も3つの入力モードをサポートしています。

配列からスカラーへ：pyarrow.Arrayを受け取り、スカラー値を返します。

• 配列のイテレータからスカラーへ：pyarrow.Arrayのイテレータを受け取り、スカラー値を返します。これは、集計スタイルの操作で大量のデータを処理する場合に役立ちます。

複数の配列のイテレータからスカラーへ：複数のpyarrow.Arrayのタプルのイテレータを受け取り、スカラー値を返します。より複雑な集計を定義できます。

Arrowテーブル関数

Arrowテーブル関数は、Arrow UDTF（ユーザー定義テーブル関数）とも呼ばれ、pyarrow.RecordBatchまたは複数のpa.Arrayを入力として受け取り、pyarrow.Tableを出力として生成します。これは、カラムナー実行を利用してPythonで実装されるテーブルイン、テーブルアウト変換の主要なパターンを表します。Arrow UDTFには以下の機能があります。

• 複数の列を返す
• ゼロ、1つ、または複数の行を生成する
• Arrow計算カーネルを利用してベクトル化されたテーブル変換を実行する

したがって、フィルタリング、行の展開、データの再構築、派生列の生成などの操作に最適です。

arrow_udtfインターフェースはシンプルさを追求して設計されており、DDL形式の文字列を使用して戻り値を定義するデコレータ構文を採用しています。この設定では、evalメソッドはPyArrowオブジェクトを入力として受け取り、PyArrow TablesまたはRecordBatchesを生成することが期待されます。このインターフェースは2つの入力モードに対応しています。テーブル引数を処理する場合、evalメソッドには、入力テーブルのすべての列をカプセル化するpa.RecordBatchオブジェクトが提供されます。

スカラー引数の場合、メソッドは各スカラー入力に対して1つのpa.Arrayオブジェクトを受け取ります。

別の例を次に示します。

このUDTFは2つの異なる方法で機能します。

Pythonでの使用法:

SQLでの使用法:

DataFrameのmapInArrowとapplyInArrowのサポート

User-Defined Functions (UDFs) および User-Defined Table Functions (UDTFs) に加えて、PySparkは、Pythonネイティブ関数をDataFrameレベルでArrowデータに直接適用できるArrow Function APIを提供します。これらのAPIは、Pandasの対応する機能 (mapInPandas、applyInPandas) と同様に動作しますが、Pandas DataFramesの代わりにpyarrow.RecordBatchとpyarrow.Tableを使用するため、PandasとArrow形式間の変換オーバーヘッドを回避できます。

• Map。DataFrame.mapInArrowは、pyarrow.RecordBatchのイテレータを別のpyarrow.RecordBatchのイテレータに変換し、フィルタリング、変換、拡張などの行レベルの操作を可能にします。

• Grouped Map。groupBy().applyInArrow()は、指定された関数を各グループに適用し、pyarrow.Tableを受け入れて返します。この機能は、データ正規化などのグループごとの変換に役立ちます。

• Co-grouped Map。 cogroup().applyInArrow()は、共通のキーに基づいて2つのDataFrameをコグループ化し、その後、各コグループに関数を適用することを可能にします。この関数は2つの pyarrow.Table入力を受け取り、単一の pyarrow.Tableを返すことが期待されます。

パフォーマンス

高コストなPandas/Arrowデータ変換を排除することで、Arrow UDFは通常、Pandas UDFよりも高速に実行され、メモリ使用量も少なくなります。2つのシンプルなUDFを比較してみましょう。

Arrow UDFはPandas UDFよりも約10%高速であり、メモリプロファイラは実行時に約40%のメモリが節約されることを示しています。

結論

Databricks Runtime 18.0では、PySparkでの高性能なPython UDF実行のために、Pandas UDFに代わるより高速で軽量なNative Arrow UDFが導入されました。Arrowデータで直接動作し、Pandas/Arrow変換のオーバーヘッドを排除することで、Arrow UDFは実行速度が約10%向上し、メモリ使用量が約40%削減され、複雑なデータ型へのサポートも強化されています。これらすべてが、使い慣れた直感的なデコレータ構文で実現されます。

さらに詳しく知りたいですか？Databricks Runtime 18.0の一部として、Databricksで今すぐNative Arrow UDFをお試しください。始めるには、既存のPandas UDFをArrow UDFに置き換えるだけです。ほとんどの場合、数行の変更で即座にパフォーマンスが向上します。完全なAPIリファレンスと追加の例については、Arrow UDFドキュメントとArrow UDTFドキュメントを参照してください。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事