Posted inPython

RustとPyO3でPythonを限界まで高速化：15分から10秒への劇的な改善

Python tutorial - IT technology blog
Python tutorial - IT technology blog

なぜPythonにはRustのような「助っ人」が必要なのか？

5GBのCSVファイルを処理したり、数百万の要素を持つ行列計算を行ったりする際に、Pythonスクリプトが遅すぎて困ったことはありませんか？私もかつて同じ状況に陥りました。最初はわずか200行のコードでスムーズに動いていたログ処理プロジェクトが、データ量の増大に伴い、実行時間が数秒から15分以上に跳ね上がってしまったのです。

その時、私には2つの選択肢がありました。1つはすべてをC++で書き直すこと（これは非常に手間がかかります）。もう1つは、Pythonのロジックはそのままに、ボトルネックとなっている部分のエンジンを「強化」することです。私は後者を選び、RustとPyO3を採用しました。

PyO3は架け橋のような役割を果たします。メモリ安全で極めて高速なRustコードを書き、それを通常のモジュールとしてPythonからimportできるライブラリとしてパッケージ化できます。これにより、Pythonの柔軟性を維持しながら、システムレベルの実行速度を手に入れることができるのです。

5分で実装：ゼロからRustモジュール作成まで

セットアップは想像以上に簡単です。わずか数コマンドで、高性能なモジュールを作成できます。

ステップ1：ツールのインストール

RustとPythonに加えて、maturinが必要です。これはRust-Pythonプロジェクトのビルドと公開に特化したツールです。

pip install maturin

ステップ2：プロジェクトの初期化

新しいディレクトリを作成し、初期化コマンドを実行します。

mkdir my_rust_module && cd my_rust_module
maturin init

バインディングの種類を聞かれたら、pyo3を選択してください。

ステップ3：Rustコードを書く

src/lib.rsを開きます。ここでは、大きなループ処理でPythonが苦手とする典型的な例として、フィボナッチ数列を計算する関数を書きます。

use pyo3::prelude::*;

#[pyfunction]
fn fibonacci_rust(n: u32) -> PyResult<u64> {
    if n <= 1 {
        return Ok(n as u64);
    }
    let mut a = 0;
    let mut b = 1;
    for _ in 0..n {
        let temp = a;
        a = b;
        b = temp + b;
    }
    Ok(a)
}

#[pymodule]
fn my_rust_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(fibonacci_rust, m)?)?;
    Ok(())
}

ステップ4：ビルドと実行

次のコマンドを入力して、Rustコードを即座にPythonモジュールとしてコンパイルします。

maturin develop

では、test.pyファイルを作成して結果を確認してみましょう。

import my_rust_module

result = my_rust_module.fibonacci_rust(90)
print(f"Rustからの結果: {result}")

PyO3の仕組みの何が特別なのか？

PyO3は、マクロを通じてPython Bindingsをスマートに処理します。

  • #[pyfunction]: Pythonから呼び出し可能なRust関数としてマークします。
  • #[pymodule]: import時に使用するモジュール名を定義します。
  • 自動型変換: PyO3はRustのu32をPythonのintに自動的にマッピングします。手動で型をキャストする必要はありません。

最大の利点は、GIL (Global Interpreter Lock)を解放できる能力です。Pythonは通常、GILのためにマルチスレッド処理が制限されますが、Rustを使えば重いタスクを真のマルチコアで実行し、その結果だけをPythonに返すことができます。

応用：Rustの構造体をPythonのクラスとして扱う

単一の関数だけでなく、データの状態を管理する必要がある場合も多いでしょう。PyO3では、RustのStruct（構造体）をPythonのClassとして扱うことができます。

#[pyclass]
struct DataProcessor {
    #[pyo3(get, set)]
    threshold: f64,
}

#[pymethods]
impl DataProcessor {
    #[new]
    fn new(threshold: f64) -> Self {
        DataProcessor { threshold }
    }

    fn process(&self, data: Vec<f64>) -> Vec<f64> {
        // しきい値より大きい要素のみを抽出
        data.into_iter().filter(|&x| x > self.threshold).collect()
    }
}

Python側での使い方は非常に自然です。

from my_rust_module import DataProcessor

proc = DataProcessor(10.5)
# 100万要素のリストを超高速で処理
result = proc.process([1.0, 15.0, 5.5, 20.0])

実践的なアドバイス：いつ使うべきか？

すべてのプロジェクトでRustを乱用すべきではありません。私の実体験に基づく注意点は以下の通りです：

  1. オーバーヘッドを考慮する: 2つの言語間のブリッジを通じたデータ転送にはリソースを消費します。Pythonで1msで終わるような処理なら、Rustで書き直す必要はありません。数百万回のループが発生するような処理に限定すべきです。
  2. ビジネスロジックはPythonに残す: 柔軟な変更が必要なロジックや全体の流れの制御はPythonに任せましょう。Rustは「計算コア」としての役割に徹するのがベストです。
  3. Cargoを活用する: Rustのエコシステムには、serde_jsonやpolarsなど、驚異的な速度を持つライブラリが豊富にあります。これらを活用し、車輪の再発明を避けましょう。

RustとPythonの組み合わせは、完璧なバランスをもたらします。Pythonの圧倒的な開発スピードと、Rustの生の処理能力を同時に手に入れることができるのです。これは、現代のビッグデータ問題を解決するための強力な武器となるでしょう。

Share:
Tags: