Python、Rust上でのデータフレームライブラリで速さ対決してみた!

※ Apache Arrow, Apache Arrow DataFusionは、米国およびその他の国におけるApache Software Foundationの登録商標または商標です。これらのマークの使用は、Apache Software Foundationによる保証を意味するものではありません。

こんにちは。5月入社のソフトウェアエンジニアのhirataと言います。日本全国地獄のように暑い日々ですが、元気で過ごしていますでしょうか？

私は入ったばかりで覚えることが多く毎日アップアップしていますが、それでも計算式を実行するエンジンを作ったり、予実の数字を分析する機能を作ったりして、楽しく過ごしております。

我々のサービス「DIGGLE」で取り扱っている「お客様の経営に関わるデータ」はそこそこ規模の大きいデータになります。このデータを高速に集計・分析できることはお客様への価値向上に直結します。今回はそれに関わる技術調査の一環で、各データフレームライブラリを比較する機会があったので記事にしてみました。

データフレームライブラリとは

データ分析のためのライブラリでPythonのpandasが有名です。2次元のdataframe、1次元のseriesというデータを格納するためのオブジェクトを持ち、データの操作を行う際にはこれらのオブジェクトにデータを格納し、オブジェクトに対して計算内容を指示していきます。

dataframe, seriesに対して一括して計算処理を行うライブラリであり、実際の計算処理はC++言語等で作成されたライブラリ内で行われるため、Python等のスクリプト言語においても大量のデータを高速に処理することができ、コンパイル言語並の高速性とスクリプト言語であることによる柔軟性とを両立できます。

このように書くと、コンパイル言語でデータフレームライブラリを使う恩恵が無いように感じられるかもしれませんが、そんなことはありません。スクリプト言語程の差では無いですが、データのメモリ配置等により大量データの計算に最適化されているため普通にプログラムして計算するよりも高速に処理できます。データ処理以外も必要な場合はコンパイル言語を使用することも選択肢となるでしょう。

ということで、今回はコンパイル言語に匹敵する計算速度が本当に出るのかの検証を、PythonとRustの二つで検証しました。

前提条件

今回の検証では以下の条件で処理時間を検証していきます。

• ファイルフォーマットはparquet形式。サンプルデータ数は100万行。
• 処理内容は三要素でgroup byしてsumする。

今回エントリーするライブラリ

今回の検証に使用するライブラリと処理系です。社内の事情で今回はPython3ではなく、Python2で計測を行います。

言語	ライブラリ
Python(2.7.17)	pandas(1.1.5) with pyarrow(6.0.1)
Python(2.7.17)	polars(0.29.0) lazy
Rust(1.69.0)	datafusion(25.0.0)
Rust(1.69.0)	polars(0.29.0) lazy

それぞれについての簡単な解説と計算を行うためのソースを以下に載せていきます。

Python pandas について

https://github.com/pandas-dev/pandas

言うまでも無いとは思いますが、データ処理ではデファクトスタンダードとなっているライブラリです。 今回はapache arrowというインメモリのフォーマットを使うためのライブラリであるpyarrowも同時に使っています。 ソースは以下のような感じです。今回は数値計算ライブラリ(blas等)等の高速化のためのライブラリ導入は行っていません。試しに入れて見たところ今回の処理内容では実行速度が変わらなかったため、素のpandasライブラリを使っています。

import time
from pyarrow import fs
import pyarrow.dataset as ds

def main():
    start_time = time.perf_counter()
    
    dataset = ds.dataset("test.parquet", format="parquet")
    dataframe = dataset.to_table(columns=['date', "account_id", "unit_id", "value"]).to_pandas()
    series = dataframe.groupby(["date", "account_id", "unit_id"])["value"].sum()
    
    print('time:', time.perf_counter() - start_time)

if __name__ == '__main__':
    main()

Python polars について

https://github.com/pola-rs/polars

こちらもapache arrowを用いたrust製のデータフレームライブラリのpythonバインディングで、かなり高速との噂です。lazyな実行とeagerな実行の両方ができます。

lazyの方が速いので今回はlazyで評価しています。今回評価を行うためのソースは以下。

import time
import polars as pl

def main():
    start_time = time.perf_counter()
    df = pl.scan_parquet("test.parquet")
    series = df.groupby(["date", "account_id", "unit_id"]).agg(pl.col("value").sum()).collect()
    
    print('time:', time.perf_counter() - start_time, series)
    print(series)

if __name__ == '__main__':
    main()

Rust dafafusion について

https://github.com/apache/arrow-datafusion

apache arrowを用いたrustのデータフレームライブラリです。こちらもApache製の様なのでapache arrowのrust向け公式ライブラリっぽいです。

今回評価を行うためのソースは以下。SQLで書いてますが、SQLじゃなくても速度は変わらずでした。

use std::time;
use std::sync::Arc;
use datafusion::prelude::*;
use datafusion::arrow::record_batch::RecordBatch;

#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
    let now = time::Instant::now();

    let ctx = SessionContext::new();

    let path = "test.parquet";
    
    // 集計用クエリの定義: SQL形式で記述可能
    let opts = ParquetReadOptions::default();
    ctx.register_parquet("facts", &path, opts).await?;
    let df = ctx.sql("SELECT date, account_id, unit_id, sum(value) FROM facts GROUP BY date, account_id, unit_id").await?;
    
    // collect() の実行で実際にクエリがexecuteされる
    let results: Vec<RecordBatch> = df.clone().collect().await?;
    
    let time = now.elapsed();
    println!("time: {:?}, result size: {}", time, df.count().await?);
    Ok(())
}

Rust polars について

https://github.com/pola-rs/polars

上記Python polarsで使用したpolarsのrustバインディングです。高速、省メモリに書けるのがrustのメリットということで、今回の最有力候補です。ソースは以下のとおり。

use polars::prelude::*;
use std::time;

#[tokio::main]
async fn main() {
    let now = time::Instant::now();
    let df = LazyFrame::scan_parquet("test.parquet", Default::default()).unwrap();
    let df = df.groupby(["date", "account_id", "unit_id"])
        .agg([col("value").sum()])
        .collect()
        .unwrap();
    let time = now.elapsed();
    println!("time: {:?}, result size: {}", time, df);
}

結果

以下計測結果です。

⁠実装	⁠実行時間
Python pandas	0.54s
Python polars	0.76s
Rust datafusion	1.53s
Rust polars	0.25s

やはり期待通りRustでのpolarsが一番速いですが、Python版でもそれほど遅くなっていませんでした。またpandasの速さが予想以上でRust polarsとの差が2倍ほどに収まっています。この測定では処理が途中でライブラリ外に出る事が無いために、pythonオブジェクトに変換するためのオーバーヘッドがほとんど無いせいでしょう。Pythonで書く容易性を考慮すると、この程度の差であれば大抵の場合はpandasで十分ということになりそうです。

またRust datafusionが期待に反して遅いですが、ここはユーザ数が膨大で歴史もあるpandasの最適化具合には敵わなかったといったところでしょうか。「今回のユースケースでは」といったところもあるかと思いますので、皆さんが実際に使用する際には事前に想定されるユースケースでの計測を行う事をお勧めします。

まとめ

pandas素晴らしいですね。気軽に書けてこれだけ速度が出れば、デファクトスタンダードになる訳だなと思いました。

また、これは反省なのですが、実はこのブログを書くにあたりメモリの使用量を取ってなかった事に気がつきました。途中で処理速度にフォーカスした為、取らなくても良いかとその時は思っていました。が、改めて見直してみると、それがどの程度だったのかがとても気になります。メモリが気になりがちな大量データの処理なのでそこを曖昧にしておかずに、きちんとメモリ使用量も計測・記録しておいた方が今後の為に良いデータとなったでしょう。次回に向けての反省点です。

We're hiring!

予実管理の明日を切り開く為に日々精進している我々と一緒に開発してくれるメンバーを大募集です。

少しでも興味があれば、ぜひ下記採用サイトからエントリーください。

herp.careers