🏹

PythonオブジェクトをRustで効率的に処理する - pyo3-arrowによるゼロコピー実装

2025/09/13に公開

FFI(Foreign Function Interface)の型変換のオーバーヘッド

RustとPythonを連携させるとき、悩ましい問題がデータの受け渡しです。

# ユーザは様々な形式でデータを渡してくる
import numpy as np
import pandas as pd

# 素数判定をRustで高速化したい
primes = rust_lib.is_prime_batch(np.array([2, 3, 4, 5, 6]))  # NumPy
primes = rust_lib.is_prime_batch([2, 3, 4, 5, 6])  # Python list
primes = rust_lib.is_prime_batch(pd.Series([2, 3, 4, 5, 6]))  # Pandas

従来のアプローチでは、これらすべてに対応するため型変換の嵐になりがちです。

// 従来のPyO3での実装: 型変換地獄
#[pyfunction]
fn is_prime_batch(py: Python, numbers: &PyAny) -> PyResult<PyObject> {
    // NumPy配列の場合
    if let Ok(numpy_array) = numbers.downcast::<PyArray1<i64>>() {
        let vec: Vec<i64> = numpy_array.to_vec()?;  // メモリコピー: 8MB (100万要素)

        let results: Vec<bool> = vec.iter()
            .map(|&n| is_prime(n))
            .collect();  // メモリアロケーション: 1MB

        // Vec → NumPy配列への変換（メモリコピー: 1MB）
        return Ok(PyArray1::from_vec(py, results).into());
    }

    // Python listの場合
    if let Ok(list) = numbers.extract::<Vec<i64>>() {  // メモリコピー: 8MB
        let results: Vec<bool> = list.iter()
            .map(|&n| is_prime(n))
            .collect();

        Ok(results.into_py(py))  // メモリコピー: 1MB
    } else {
        Err(PyTypeError::new_err("Type conversion required"))
    }
}

各変換でメモリコピーが発生し、パフォーマンスが悪化します。
本記事では、Apache Arrowを中心に据えることで、この問題を効果的に解決する方法を紹介します。

Arrow-nativeアプローチによる型変換の排除

従来のFFI設計の問題

典型的なRust-Python FFIのデータフローを示します。

100万要素の配列では、入出力時にデータサイズに比例したメモリコピーが発生します（int64なら各8MB）。

Arrow-native設計の仕組み

Apache Arrowを中心に据えると以下のようになります。

型変換が不要になり、メモリコピーはゼロになります。

実装例 - 素数判定を高速化

ここでは、素数判定を例に従来のアプローチとArrow-nativeアプローチの違いを具体的に確認します。

NumPyのufuncやベクトル化した関数は非常に高速です。しかし、素数判定のような数式化できず、条件判定や繰り返しが要求される処理では、Pythonのオーバーヘッドが顕著に現れます。
このような処理はRustによる高速化の恩恵を受けられます。

Rust側の実装

従来のPyO3実装（型変換あり）

従来の実装では次のような課題があります。

入力時: NumPy配列やリストが自動的に Vec<i64> にコピーされる
出力時: Vec<bool> がPythonオブジェクトにコピーされる

// src/lib.rs
use pyo3::prelude::*;

/// 素数判定（単純な実装）
fn is_prime(n: i64) -> bool {
    if n < 2 { return false; }
    if n == 2 { return true; }
    if n % 2 == 0 { return false; }

    let sqrt_n = (n as f64).sqrt() as i64;
    for i in (3..=sqrt_n).step_by(2) {
        if n % i == 0 { return false; }
    }
    true
}

// 従来の実装: Vec型で受け取る
#[pyfunction]
fn is_prime_batch_old(numbers: Vec<i64>) -> PyResult<Vec<bool>> {
    // Python → Vec<i64>（メモリコピー発生）
    let results: Vec<bool> = numbers.iter()
        .map(|&n| is_prime(n))
        .collect();

    Ok(results)  // Vec<bool> → Python（メモリコピー発生）
}

Arrow-native実装（ゼロコピー）

Arrow-nativeによって次の効果が期待できます。

入力: Arrow配列のメモリを直接参照（コピーなし）
出力: Arrow配列として直接返却（コピーなし）
メモリ効率が大幅に向上

// src/lib.rs
use pyo3::prelude::*;
use pyo3_arrow::{PyArray, PyArrowResult};
use arrow::array::{Int64Array, BooleanArray};
use std::sync::Arc;

#[pyfunction]
fn is_prime_batch(py: Python, numbers: PyArray) -> PyArrowResult<PyObject> {
    // Arrow配列を直接受け取る（ゼロコピー）
    let array = numbers.as_ref();

    // Int64配列として解釈
    let int_array = array
        .as_any()
        .downcast_ref::<Int64Array>()
        .ok_or_else(|| PyErr::new::<pyo3::exceptions::PyTypeError, _>(
            "Expected Int64Array"
        ))?;

    // 各要素に素数判定を適用
    let results: BooleanArray = int_array
        .iter()
        .map(|n| n.map(is_prime))
        .collect();

    // 結果もArrow配列として返す（ゼロコピー）
    Ok(PyArray::from_array_ref(Arc::new(results)).to_arro3(py)?)
}

Python側の使用例

Pythonユーザからは、Rustの実装を意識することなく、次のようなコードで利用できます。

import pyarrow as pa
import numpy as np
from my_rust_lib import is_prime_batch

# PyArrow配列で渡す
numbers = pa.array([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
result = is_prime_batch(numbers)
print(result)  # [true, true, false, true, false, true, false, false, false, true]

# NumPy配列もそのまま動く
np_numbers = np.array([2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
result = is_prime_batch(np_numbers)  # 自動的にArrowに変換される！

NumPyでも動作

前述のコードでは、 pyarrow.array だけでなく、 numpy.array も動作します。これはpyo3-arrowの機能によって実現されています。

PyArrayのextractメソッドが呼ばれる
NumPy配列を検出
Buffer Protocol経由でゼロコピーアクセス
Arrow配列として処理

これによって、ユーザが型変換コードを書くことなく、NumPyの配列を効率的に処理できます。

パフォーマンス測定

すべての実装で同一アルゴリズム（エラトステネスの篩）を使用し、性能を比較しました。

ソースコード: https://github.com/drillan/prime-arrow-example

100万要素での性能比較

実装方法	実行時間	Python比	説明
純粋Python	3,267ms	1.0x	基準実装
NumPy最適化	23ms	142x	スライシング最適化版
pyo3-arrow (Zero-Copy)	3.5ms	942x	完全ゼロコピー実装

100万要素の素数判定における処理時間比較（対数スケール） - pyo3-arrowは純粋Python比で942倍高速

import time
import numpy as np
from prime_arrow_example import get_primes_up_to

# テストデータ: 100万までの素数を取得
limit = 1_000_000

# pyo3-arrow実装（ゼロコピー）
start = time.perf_counter()
primes = get_primes_up_to(limit)  # Arrow配列として返る
arrow_time = time.perf_counter() - start

print(f"pyo3-arrow: {arrow_time*1000:.1f}ms")
print(f"素数の個数: {len(primes)}")

スケーラビリティ - データサイズ別性能

データサイズ	純粋Python	NumPy	pyo3-arrow	高速化率
10,000	1.6ms	0.08ms	0.03ms	53x
100,000	174.4ms	1.2ms	0.3ms	581x
1,000,000	3,267ms	23ms	3.5ms	942x

データサイズ別の処理時間推移（対数スケール） - データサイズが増えるほどpyo3-arrowの優位性が増大

データサイズが大きくなるほど、Arrow-nativeアプローチの優位性が顕著になります。

メモリ効率 - ゼロコピーの実証

Arrow-nativeの実測値

import tracemalloc
import numpy as np
from prime_arrow_example import is_prime_batch

# 100万要素のテストデータ
numbers = np.arange(1, 1_000_001)

# Arrow-native実装のメモリ使用量
tracemalloc.start()
result = is_prime_batch(numbers)  # ゼロコピー処理
current, peak = tracemalloc.get_traced_memory()
print(f"追加メモリ使用量: {peak / 1024 / 1024:.1f}MB")
tracemalloc.stop()

実行結果

追加メモリ使用量: 0.1MB

100万要素処理時の追加メモリ使用量 - pyo3-arrowはゼロコピーにより追加メモリがほぼゼロ

NumPy配列のメモリを直接参照し、結果もゼロコピーで返却するため、追加のメモリ使用量はほぼゼロです。

Unsafeコードの隠蔽とメモリ安全性

自前実装の危険性と安全性の課題

Arrow C Data Interfaceを直接扱うと、以下のような重大な安全性リスクが発生します。

// 危険なunsafeコード
unsafe {
    let schema = Box::from_raw(c_schema as *mut arrow::ffi::FFI_ArrowSchema);
    let array = Box::from_raw(c_array as *mut arrow::ffi::FFI_ArrowArray);

    let data = arrow::array::make_array(
        arrow::array::ArrayData::try_from_ffi(schema, array)?
    );

    // メモリ管理を間違えるとクラッシュ...
}

潜在的なセキュリティリスク

Use-After-Free: 解放済みメモリへのアクセス
Double-Free: 同じメモリを2回解放
メモリリーク: 適切に解放されないメモリ
バッファオーバーフロー: 境界チェックなしのメモリアクセス
データ競合: 複数スレッドからの同時アクセス

メモリ安全性の保証

Arrow-nativeアプローチは、Rustの所有権システムとArrowの設計により、上記のリスクを構造的に排除します。

1. 所有権の明確な管理

// Arrow配列の所有権は明確に管理される
fn safe_processing(py: Python, data: PyArray) -> PyArrowResult<PyObject> {
    // `PyArray`がPython側の参照を保持
    let array = data.as_ref();  // 借用のみ、所有権は移動しない

    // 新しい配列を作成（元のデータは変更しない）
    let result = process_immutably(array);

    // 新しい所有権を持つ配列を返す
    Ok(PyArray::from_array_ref(Arc::new(result)).to_arro3(py)?)
}

2. ライフタイムの自動管理

// ライフタイムによるメモリ安全性の保証
impl PyArray {
    // `'a` は`PyArray`の生存期間と紐付けられる
    pub fn as_ref<'a>(&'a self) -> &'a dyn Array {
        // `self`が生きている限り、返される参照も有効
        &self.inner
    }
}

3. データ競合の防止

Arrow配列は基本的にイミュータブルであり、データ競合が構造的に防止されます。

// 並列処理でも安全
use rayon::prelude::*;

fn parallel_processing(array: &Int64Array) -> Vec<bool> {
    array.values()
        .par_iter()  // 並列イテレーション
        .map(|&n| is_prime(n))  // 読み取りのみ、競合なし
        .collect()
}

これらの安全性メカニズムを手動で実装するのは複雑でエラーが起きやすいため、pyo3-arrowがすべてのunsafeコードを隠蔽し、安全なAPIを提供しています。

pyo3-arrowによる安全な抽象化

// 安全で簡潔
fn process(data: PyArray) -> PyArrowResult<PyObject> {
    let array = data.as_ref();  // すべての`unsafe`が隠蔽されている
    // 処理...
    Ok(PyArray::from_array_ref(result).to_arro3(py)?)
}

実プロジェクト(QuantForge)の事例

最後に、筆者が開発している金融計算ライブラリQuantForgeでの実例を紹介します。

Before: 856行の型変換層

# python/quantforge/numpy_compat.py (294行)
class NumpyCompatibilityLayer:
    def convert_to_arrow(self, data):
        if isinstance(data, np.ndarray):
            # 複雑な変換ロジック...
            
# python/quantforge/wrappers.py (177行)
class BatchWrapper:
    def prepare_inputs(self, *args):
        # 各引数を検査して変換...

# python/quantforge/arrow_api.py (335行)
class ArrowAPIAdapter:
    # さらなる変換層...

After:

// bindings/python/src/utils.rs
pub fn pyany_to_arrow(_py: Python, value: &Bound<'_, PyAny>) -> PyResult<PyArray> {
    // 1. If already an Arrow array, check if it needs conversion to Float64
    if let Ok(array) = value.extract::<PyArray>() {
        let array_ref = array.as_ref();

        // Check if the array is already Float64
        if array_ref.data_type() == &arrow::datatypes::DataType::Float64 {
            return Ok(array);
        }

        // If it's Int64, convert to Float64
        if array_ref.data_type() == &arrow::datatypes::DataType::Int64 {
            use arrow::compute::cast;

            let casted = cast(array_ref, &arrow::datatypes::DataType::Float64).map_err(|e| {
                PyValueError::new_err(format!("Failed to cast array to Float64: {e}"))
            })?;
            let array_ref = casted;
            return Ok(PyArray::from_array_ref(array_ref));
        }

        // For other types, try to cast
        use arrow::compute::cast;
        let casted = cast(array_ref, &arrow::datatypes::DataType::Float64)
            .map_err(|e| PyValueError::new_err(format!("Failed to cast array to Float64: {e}")))?;
        let array_ref = casted;
        return Ok(PyArray::from_array_ref(array_ref));
    }

    // NumPy配列やPythonリストの処理は省略（実際の実装はリポジトリ参照）
    // ...
}

結果として8-9割のコードを削減しています。

まとめ

本記事では、RustとPythonのFFIにおける型変換のオーバーヘッドを、pyo3-arrowで解決する方法を紹介しました。

主要な成果

パフォーマンス:

高速化を達成
完全なゼロコピー実装により追加メモリ使用量をほぼゼロに削減
データサイズが大きくなるほど優位性が増大

開発効率:

型変換コードを8-9割削減（QuantForgeプロジェクトでは856行を削除）
NumPy、PyArrow、Pandasなど複数の入力形式を統一的に処理
unsafeコードの完全な隠蔽により安全性を確保

pyo3-arrowが適している場面

大規模データ処理: 数百万〜数億要素のデータを扱うケース
リアルタイム処理: レイテンシーが重要な金融計算や科学計算
メモリ制約環境: メモリコピーを避けたい組み込みシステムやエッジデバイス
Parquetファイル処理: Arrow形式と直接互換性があるParquetファイルの高速読み書き（ゼロコピーで処理可能）
既存のArrowエコシステムとの統合: Polars、DuckDBなどとの連携
NumPyでベクトル化できない処理: 複雑な条件分岐、可変長ループ、再帰的処理など

実装時の留意点

Arrow形式への統一: 入出力をArrow形式に統一することで最大の効果を発揮
Buffer Protocolの活用: NumPy配列は自動的にゼロコピー変換される
戻り値の型変換: 戻り値はArrow配列のため、NumPyで後続処理する場合は.to_numpy()などの変換が必要
エラーハンドリング: 型の不一致時の適切なエラーメッセージの実装