PyO3 を使って Python 内で Rust ライブラリを実行する

PyO3を使ってPythonでRustライブラリを実行する方法

Pydantic v2のようなデータ検証ライブラリは、Pythonアプリケーションで広く利用されていますが、そのコア部分はRustで書かれており、PyO3というツールチェインを通じてPythonから呼び出されています。この記事では、このPyO3を使って、Rustで書かれたJSONパーサーをPythonからインポート可能なパッケージとして公開する手順を解説します。

最後のステップである、Rustの結果をPythonオブジェクトに変換する処理は、特に理解しておくべき重要な点です。このようなパーサーの場合、この変換コストが解析処理自体のコストを上回ることがあります。

Rustからインポートまでの4つのステップ

RustコードをPythonで利用可能にするには、以下の4つのステップが必要です。

1. 通常のRustモジュールを作成する。
2. PyO3マクロでアノテーションを付ける。
3. Maturinでコンパイルしてインストールする。
4. 結果をインポートする。

［画像: Rust in Python: write a Rust module, annotate it with PyO3 macros, build it with maturin, import the shared library］

#[pyfunction]と#[pymodule]は、この連携を行う2つのRustマクロです。Rustの属性マクロはPythonのデコレータに似ており、対象となる関数を書き換えて、Pythonからの呼び出しを可能にし、境界での型変換や参照カウントを処理します。

Maturinは、クレートを共有ライブラリ（.so、.dylib、.dll）としてコンパイルし、仮想環境に配置するため、import文で簡単に利用できるようになります。このセットアップ全体の手順については、How to run Rust in Python with PyO3 and Maturinで詳しく解説しています。

最初のチュートリアルでは単一の数値を返しましたが、今回はスカラー値ではなく構造体を返す場合に何が起こるか、というより興味深い点から話を始めます。

パーサーはまずRustの値を出力する

このパーサーが出力する構造体はJSONツリーであり、この記事の残りの部分で実行例として使用します。私たちのPython to Rust cohortでは、開発者が6週間かけてRustでJSONパーサーを一から書き、serdeを使わない手動のトークナイザーと再帰下降パーサーを作成し、それをPyO3経由でPythonに公開します。

ここで紹介するコードは私の実装です。

パーサーはプレーンなRustのenumを出力します。Rustのenumは複数の形状のうちの1つを持ち、各バリアントはデータを保持できるため、JSONツリーをきれいにマッピングできます。

pub enum JsonValue {
    Null,
    Boolean(bool),
    Number(f64),
    String(String),
    Array(Vec<JsonValue>),
    Object(HashMap<String, JsonValue>),
}

このツリーは完全にRust内に存在し、Pythonからは見えません。PyO3レイヤーは、その上に薄いアダプターとして機能します。

1つの関数の公開

関数をPythonに公開するには、2行のコードが必要です。

#[pyfunction]
fn parse_json<'py>(py: Python<'py>, input: &str) -> PyResult<Bound<'py, PyAny>> {
    parse(input)?.into_pyobject(py)
}

Pythonの読者にとって、このシグネチャが最も興味深い部分でしょう。

• py: Python<'py>はPythonインタプリタへのアクセスを表すトークンであり、Pythonオブジェクトにアクセスする必要があるPyO3 APIに渡すものです。従来のPythonビルドでは、このアクセスはGIL（Global Interpreter Lock）を保持することに関連付けられています。PyO3はこれをあなたに渡し、Pythonオブジェクトに触れる場所ならどこでも渡す必要があります。
• Bound<'py, PyAny>は、任意の型のPythonオブジェクトへのハンドルであり、Rust側から見るとPyObjectに相当します。
• PyResult<T>はResult<T, PyErr>であり、値を返すか、PyO3がPython例外として発生させるエラーを返します。
• ?演算子はエラーを伝播させます。parseが失敗した場合、関数は早期リターンし、Pythonは例外を受け取ります。そうでなければ、JsonValueをアンラップして続行します。

したがって、parse(input)?が実際の処理を行い、.into_pyobject(py)が呼び出し元が要求したPythonオブジェクトを構築します。この最後の呼び出しがコストの発生源です。ツリー内のノードごとにPythonオブジェクトを作成する必要があり、大きなドキュメントでは、解析処理自体よりも多くの作業が発生する可能性があります。

戻り値の変換がコストのかかる部分

この変換が無料ではない理由はここにあります。.into_pyobjectはJsonValueツリー全体をウォークし、ネイティブPythonオブジェクトとして再構築します。オブジェクトごとにdict、配列ごとにlist、リーフごとにfloatまたはstrです。この翻訳は、Rustの値をPythonの値に変換するためにPyO3が呼び出すIntoPyObjectトレイトを実装することで提供されます。

impl<'py> IntoPyObject<'py> for JsonValue {
    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            JsonValue::Null => Ok(py.None().into_bound(py)),
            JsonValue::Number(n) => Ok(n.into_pyobject(py)?.to_owned().into_any()),
            JsonValue::Object(obj) => {
                let py_dict = PyDict::new(py);
                for (k, v) in obj {
                    py_dict.set_item(k, v.into_pyobject(py)?)?;  // 再帰呼び出し
                }
                Ok(py_dict.into_any())
            }
            // ... 配列、文字列、ブール値
        }
    }
}

100,000個の値を持つドキュメントの場合、境界で100,000個程度のPythonオブジェクトが作成されることになります。これはすべて、解析が完了した後に行われます。大きなドキュメントでは、このマテリアライゼーションループが、解析処理ではなく、エンドツーエンドの時間を支配する可能性があります。

エラーも同様の方法で境界を越える

戻り値だけが変換を必要とするわけではありません。解析の失敗は型付けされたRustのエラーですが、Pythonは例外を期待しています。Rustが型を別の型に変換するために使用するFrom実装（トレイト）の1つにより、?演算子で処理できます。

impl From<JsonError> for PyErr {
    fn from(err: JsonError) -> PyErr {
        match err {
            JsonError::UnterminatedString { position } => PyValueError::new_err(
                format!("Unterminated string starting at position {position}")
            ),
            // ...エラーバリアントごとに1つのアーム、位置を保持
        }
    }
}

これにより、不正な入力はオフセットを保持したValueErrorとして発生します。ファイル読み取りパスも同様に無料で扱われます。std::io::Errorは既に一致するPython例外に変換されるため、パスが見つからない場合はFileNotFoundErrorが発生します。

呼び出し元は、Rustレイヤーが漏洩することなく、Pythonのセマンティクスを受け取ります。

独自のポートに何を意味するか

ポートするRust関数がスカラー値を返す場合は、それをポートして次に進んでください。境界は通常、無視できるほど小さいです。

大きな構造体を返す場合、変換が実際のコストとなり、パーサー自体が高速になったら次に最適化すべき点です。PyDictの事前割り当ては、わずかに役立つかもしれませんが、より大きな利点はアーキテクチャにあります。呼び出し元がすべてにアクセスしない場合は、ツリー全体をマテリアライズしないでください。遅延ロードされるRustベースのビューを返し、オンデマンドでPythonオブジェクトを構築します。

したがって、PyO3を利用する際は、アルゴリズムだけでなく、境界のパフォーマンスをプロファイルしてください。Rustを高速に実行させることは簡単な半分です。出口で構築するもの、つまりRustの値からPythonオブジェクトへの旅が、ポートが価値があったかどうかを決定する半分です。

AIはコンパスではなくアクセラレータです。私は開発者がより速く出荷し、出荷したものを理解できるように指導しています。パスを見る →