Rust でテキストファイルからデータを読み出す Python 拡張を書いたら爆速だった話

普段、実験をしていて得られるデータを np.loadtxt で読んでいるが、ためしにこの部分を Rust で書いて Python から呼び出してみたら速かったという話を書きます。

環境
準備
コード
速度比較
まとめ

最近こんな記事を書きました。
pyhaya.hatenablog.com
この記事では、Rust でテキストデータを読み込んだときにどれくらいの速度が出るのだろうということを試してみた記事です。この中で、最後に Python の numpy との比較をしているのですが、ここで結構 Rust と numpy の間に差ができており、Rust で Python 拡張を書いたときにどれくらいスピードが出るのか気になりました。

Rust から Python 拡張を書くのはそんなに難しくないので、試してみた結果をここに書くことにしました。

環境

Ubuntu 19.10
Cargo 1.39.0-nightly (<- stable だと動かないので注意)
Python 3.7.3

rustup を入れていれば nightly は以下の記事のように簡単にインストールすることができます。
qiita.com

準備

適当なディレクトリで次のコマンドを打ちます。("speed" はディレクトリの名前）

cargo new --lib speed

生成されたディレクトリ内の Cargo.toml を修正する。

[package]
name = "speed"
version = "0.1.0"
authors = ["Your name <mail address>"]
edition = "2018"

[lib]
name="speed"
crate-type=["cdylib"]

[dependencies]

[dependencies.pyo3]
version="0.8"
features=["extension-module"]

コード

Numpy と同じく、loadtxt() という名前の関数を定義して Rust で書き下します。このとき、返り値は PyResult型である必要があることに注意します。

src/lib.rs

use pyo3::prelude::*;
use std::fs::read_to_string;

#[pymodule(speed)]
fn loadtxt(py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "loadtxt")]
    fn loadtxt(_py: Python, filename: &str) -> PyResult<(Vec<f32>, Vec<f32>)> {
        let data = read_to_string(filename);
        let xy = match data {
            Ok(content) => content,
            Err(error) => { panic!("Could not open file: {}", error); }
        };

        let xy_pairs: Vec<&str> = xy.trim().split("\n").collect();
        let mut x: Vec<f32> = Vec::new();
        let mut y: Vec<f32> = Vec::new();

        for pair in xy_pairs {
            let p: Vec<&str> = pair.trim().split(" ").collect();

            x.push(p[0].parse().unwrap());
            y.push(p[1].parse().unwrap());
        }

        Ok((x, y))
    }

    Ok(())
}

これをビルドすれば、Python側から呼び出せるバイナリデータが出来上がります。

cargo build --release

速度比較

target/release 内に含まれる libspeed.so を speed.so に名前を変えてわかりやすいように別のディレクトリにコピーします。そのディレクトリで Python ファイルを作って速度を比較してみました。

まず、読み出すファイルを作成します。999999行のデータです。

data = ""
num = 1000000
for i in range(num):
    data += f"{i} {i*2}"
    if i < num - 1:
        data += "\n"
with open("example.txt", "w") as f:
    f.write(data)

では、実際に速度を比較してみます。

from speed import loadtxt
import numpy as np
import time

if __name__ == "__main__":
    filename = "example.txt"
    
    # Rust
    t1 = time.perf_counter()
    x, y = loadtxt(filename)
    print("Rust time: ", time.perf_counter() - t1)

   # Numpy
    t2 = time.perf_counter()
    x, y = np.loadtxt(filename).T
    print("Numpy time: ", time.perf_counter() - t2)

Rust time:  0.29457211199951416
Numpy time:  3.338321666000411

Rust で書いたほうは Numpy よりも 10 倍以上速いという結果になりました。

まとめ

この記事では、テキストデータを Python で読み出すためにNumpyのloadtxtを使った場合と、Rustで自作の拡張を作った場合の２通りのコードを書き、速度の比較を行いました。その結果、Rustのほうが 10 倍以上速いという結果になりました。ただ、Rust の方は読み出すデータが２列からなるということを知った上で書いてあるので、そのせいで Numpy より早くなっているという可能性はあるかも知れません。今後はそこらへんも検証していきたいと思っています。