Rust式プラグインによるPolarsの拡張

fenicはセマンティックDataFrameライブラリです。散乱した非構造化データ上でAIおよびLLMパイプラインを構築するためのPySparkスタイルのAPIを備えています。そのローカルエンジンはPolarsです。この記事は、それを構築する際に直面した特定の問題と、それを解決したPolarsの機能、すなわち式プラグインシステムについて述べています。最終的に、Polarsの式エンジンを拡張する9つのRustプラグインを作成しました。以下に、その理由と、実際のコードを用いた構築方法を説明します。要約：AIパイプラインがテキストに対して必要とする操作（チャンキング、プロンプトテンプレート、jq、ファジーマッチング、Markdownおよびトランスクリプト解析、リッチな型キャスト）はPolarsにはありません。これらをPython UDFとして実行すると遅く、コンポジションが壊れます。RustでPolarsの式プラグインとしてこれらを記述すると、ネイティブな式になります。これらはArrow上でインエンジンで実行され、宣言された型を維持し、単一の式ツリー内で組み込みオペレーションとコンポーズできます。UDFとプラグインのどちらを選ぶか迷っているなら、プラグインを選ぶべきです。 なぜこれを構築したのか fenicとは何か fenicは、PySparkにモデル化されたAPIを持つ、AIおよびLLMパイプライン構築のためのDataFrameライブラリです。Polarsの経験者にとって重要なのは、Polarsがfenicのコア実行エンジンであるということです。記述するすべてのDataFrame操作は、Polars式、またはpl.DataFrames上のプランになります。セマンティックオペレーター、LLMのmap/extract/classify、埋め込み、類似性結合などは、すべて同じメカニズムの上に成り立っています。したがって、実際には、fenicができることは、Polarsで表現できることに限定されます。これは意図的な選択です。Polarsは、高速でカラム指向、Arrowネイティブなエンジンであり、実際の式言語と遅延評価のオプティマイザーを提供します。私たちはそのどれも再構築するつもりはありませんでした。私たちはそれに付加価値を与えたかったのです。 やりたかったこと fenicがターゲットとするワークロードは、非構造化テキスト上のパイプラインです。ドキュメント、チャットログ、トランスクリプト、スクレイピングされたJSON、Markdownなどです。具体的には、Polarsにネイティブな同等物がない行レベルの操作が必要です。 * 埋め込みや検索のために、トークン数でサイズが決まるオーバーラップするウィンドウにドキュメントをチャンク化する。 * 行ごとにプロンプトテンプレートをレンダリングする。LLM呼び出しの前に、列の構造体を変数として使用する実際のJinja。 * JSON列をjqでクエリし、Markdownを構造化されたASTに解析する。 * 重複排除と結合のために、文字列をファジーマッチングする（6つの編集距離メトリック）。 * トランスクリプト（SRT/WebVTT）を型付けされたタイムスタンプ付きキューレコードに解析する。 * Polarsの物理的なデータ型では直接モデル化できない、fenicのリッチな論理型（埋め込み、Markdown、型付き構造体）に値をキャストする。 これらの各操作は、列全体で実行され、型付けされた結果を生成し、より大きなDataFrameパイプラインの中間にスロットインする必要があります。脇に置かれるのではなく。 明白なアプローチが不十分だった理由 Polarsにカスタム操作を追加する最も一般的な方法は、Python UDFです。行ごとの作業にはmap_elements、Series全体にはmap_batchesを使用します。真に不透明でI/Oバウンドな作業（LLM API呼び出しなど）には、依然として適切なツールであり、fenicはまさにそのためにmap_batchesを使用しています。しかし、上記のテキスト操作では、UDFは3つの点でコストが高すぎました。 * 速度。 map_elementsは、GILの下で、各値に対してPythonオブジェクトのラウンドトリップを伴って、行ごとにPythonを実行します。数百万行のトークン化やファジーマッチングでは、それがボトルネックになります。作業自体ではなく。 * コンポジション。 これが実際に問題でした。Polarsが高速なのは、式ツリー全体を計画および実行するためです。あるステップが不透明なPythonコールバックである瞬間、エンジンはそれを透過できません。最適化の障壁となり、マテリアライゼーションを強制し、単一パスパイプラインを壊します。3つのUDFをチェーンすると、エンジンとの間で3回の往復が発生します。 * 型。 UDFの出力型は、緩やかにアサートして期待するものです。List<String>、Struct、固定サイズの埋め込み配列など、実際の宣言されたデータ型を生成する操作が必要でした。これにより、残りのプランはそれらに対して型チェックできます。 代替案はさらに悪かったです。Polarsをフォークしてネイティブカーネルを追加することは、永遠にフォークを所有することを意味します。DataFrameの外でテキスト作業を行い、前処理してからロードすることは、Polarsを使用する理由そのものである遅延評価とコンポジションを捨ててしまいます。これらの操作を式エンジンのファーストクラスの市民にしたかったのです。それらの隣ではなく。 なぜプラグインなのか Polarsには、まさにこのための専用の回答があります：式プラグインです。Rustでカーネルを記述し、登録すると、通常のpl.Exprになります。エンジンにとっては、組み込みのものと区別がつきません。これは、UDFルートで欠けていたすべてのボックスをチェックします。エンジン内で、Rustで、Arrowバッファ上で実行されます。ベクトル化され、並列化可能で、ストリーミングの対象となり、Pythonや行ごとのオブジェクトラウンドトリップはありません。pl.Exprを返すため、コンポーズできます。プラグインは互いに、またネイティブオペレーションと、エンジンが全体として計画する単一の式ツリーでチェーンされます。出力データ型を宣言するため、プランはカスタムステップを通じて型付けされたままになります。Rustは、ホットループ（jaq、minijinja、rapidfuzz、tiktoken、Markdownパーサー）のための成熟したエコシステムを提供し、何も再実装する必要はありません。 私たちが定めたルールは、「すべてをRustで書き直す」ではありません。ネイティブPolarsは高速パスのままです。プラグインは、Polarsが表現できないギャップを埋めるだけです：トークナイザー、jqエンジン、キャプチャグループインデックス自体が列である正規表現。ファジーマッチングでさえ、6つのプリミティブカーネルのみをRustに保持し、通常のPolars式から高階比率をコンポーズします。pyo3-polarsは、FFI、Arrowマーシャリング、およびキーワード引数ブリッジを生成するため、記述コストは低いです。私たちは9つ書きました。 何が得られたか 最終状態は、メカニズムの前に述べる価値があります。fenicのすべてのテキスト操作は、現在ネイティブなPolars式です。これらは共有Arrowメモリ上でエンジン内で実行され、型を維持し、そして最も重要な部分です。これらは、Pythonのラウンドトリップゼロで、単一の式でネイティブPolarsオペレーションとコンポーズされます。Markdownドキュメントを解析し、そのASTをjqでフィルタリングし、結果にインデックスを付け、型付き構造体にキャストすることは、1つの式です。Polarsはそれを単一パスで計画および実行し、カスタムRustと組み込みリストオペレーションを並べて実行します。ピースが揃ったら、その正確な例に戻ります。この記事の残りは、Python登録からArrowメモリまで、実際のコードのみを使用して、どのように構築されるかです。 構築方法：実装ウォークスルー メンタルモデル：契約の周りの2つの薄いレイヤー Polars式プラグインは、2つの小さなコード片と、それらの間の明確に定義された契約です。 Python側。 関数を登録して、ネイティブPolarsのように見せる：expr.my_namespace.my_op(...). Rust側。 &[Series] を受け取り、PolarsResult<Series> を返し、出力データ型を宣言する関数。 それらの間のすべては、pyo3-polarsによって生成されます。Arrowを介してSeriesをFFI境界を移動させ、キーワード引数をマーシャリングし、シンボルルックアップを配線します。FFIを手動で触ることはありません。 fenicのjson.jqオペレーターのPythonサーフェス全体は次のとおりです。 ```python # src/fenic/_backends/local/polars_plugins/json.py from pathlib import Path import polars as pl from polars.plugins import register_plugin_function PLUGIN_PATH = Path(__file__).parents[3] @pl.api.register_expr_namespace("json") class Json: """Namespace for JSON-related operations on Polars expressions.""" def __init__(self, expr: pl.Expr) -> None: self.expr = expr def jq(self, query: str) -> pl.Expr: return register_plugin_function( plugin_path=PLUGIN_PATH, function_name="jq_expr", args=self.expr, kwargs={"query": query}, is_elementwise=True, ) ``` 作業を行っているのは2つのPolars APIです。 @pl.api.register_expr_namespace("json") は、プロセス内のすべてのpl.Exprに.jsonアクセサをボルトで固定します。インポート後、pl.col("payload").json.jq(".name") は、Polars式が合法な場所ならどこでも合法な式です。 register_plugin_function(...) は、エンジンが評価したときに、plugin_pathにあるコンパイル済みライブラリをdlopenし、function_nameという名前のシンボルを探し、入力Seriesを渡し、結果を読み戻す通常のpl.Exprを返します。それが全体のアイデアです。.jsonは特別なものではありません。