Pythonから利用できるRust製超高速データ分析ライブラリPolarsの実力：Rustで始めるWebアプリケーション（終）

RustでWebアプリケーションを開発する際に基礎となる要素技術からRustの応用まで、Rustに関するあれこれを解説する本連載。第3回は、Rust製の高速データ分析ライブラリであるPolarsの速度を簡易的に検証し、考察する。

[藤田直己，paiza]

　paizaでWebエンジニアをやっています藤田と申します。前回の連載では、RustでWebアプリの基礎となるセッション管理と、SNSのAPIサーバを構築するための実装概略、Rustの強力な型システムによるサーバサイドアプリケーションの記述性について示しました。

　今回は、趣向を変えてRust製の高速データ分析ライブラリである「Polars」を利用し、その速度を簡易的に検証、考察します。今回のプロジェクトもGitHubのサンプルリポジトリを用意していますので、コードを実行する際はご利用ください。

Polarsとpandas

　PolarsはPythonの小規模データ分析文脈でよく用いられるpandasを強く意識したライブラリです。pandasがパンダならば、それに対してPolarsはホッキョクグマというわけです。

　どちらも「データフレーム」とよばれる抽象データ型が使いやすいインタフェースを形成しており「高速にデータ処理できるExcelのような分析ライブラリ」あるいは「インメモリリレーショナルデータベース」という感じで、気軽に扱えるツールとなっています。

　利用可能なデータソースとして、リレーショナルデータベースはもちろん、csvや「Microsoft Excel」「Apache Parquet」などのファイル、また「Amazon S3」などにあるデータを扱うことができる点も共通しています。

　相違点として、pandasがNumPy配列というC言語の構造体的なシンプルなデータ構造をバックエンドとして持つ一方、PolarsはApache Arrow memory modelという列指向データ構造を有しており、分析クエリに適しています。またPolarsの公式サイトではpandas含む他のデータ分析ライブラリよりも卓越した処理速度を持っていることが示されています。

　さらに、Polarsにはpandasにない遅延評価や並列処理などクエリ最適化の機能が織り込まれており、計算資源を有効活用できます。一方でクエリ最適化を行えるようにする都合上pandasとはAPIが異なる部分があります（どちらかといえば大規模データ分析ライブラリの「Apache Spark」にややAPIが似ています）。そしてpandasはPythonのAPIしか提供されていませんが、PolarsはPythonおよびRustのAPIが提供されています。

　探索的データ分析（Exploratory Data Analysis、EDA）をあえてRust APIで実施するといったユースケースは考えにくいため、Pythonで作成した分析、学習結果をRust製アプリケーションにシームレスに組み込むなどのユースケースが考えられるでしょう。今回はRustからの利用についても試してみます。

今回の試験の実行環境

　パフォーマンスは実行環境に依存しますが、参考までに筆者のプログラムの実行環境を下記に示します。物理コアが16個あるので、並列処理可能だと大きく高速化できます。

• OS：Debian 11（Linux 5.10.0-17-amd64）
• CPU：AMD Ryzen 9 3950X 16-Core Processor
• RAM：G.Skill F4-3200C16-32GVK×4（DDR4-3200 32GB×4）
• SSD：Crucial CT1000MX500SSD1（1000GB Serial ATA 6Gb/s）

Hello Polars

　まずは以前の記事でも取り上げたアヤメのデータセットについて典型的な集計クエリを実行します。

　アヤメの種類（class）ごとに、がく片の長さ（sepal length）、がく片の幅（sepal width）、花弁の長さ（petal length）、花弁の幅（petal witdth）の平均および標準偏差を計算します。がく片の長さの平均値の降順でデータを並び替えることにします。またPythonの実行時間計測ツール「timeit」を用いてpandasとPolarsの実行速度を比較します。

　アヤメのデータセットはもともと150行の非常に小さなデータセットであるので、実践ではより巨大なデータを処理することをイメージし、1万回データを繰り返したもの（つまり150万行のデータ）を対象として処理時間を計測します。Pythonおよび今回利用するライブラリのバージョンは以下のようになります。表示には、バージョン管理ツールPoetryを利用しました。

# pyproject.toml
[tool.poetry]
authors = ["NaokiFujita"]
description = "test script for performance comparison between polars and pandas"
name = "web_engineer_in_rust"
version = "0.1.0"
[tool.poetry.dependencies]
pandas = "^1.5.0"
python = "^3.9"
polars = "^0.14.18"
[tool.poetry.group.dev.dependencies]
autopep8 = "^1.7.0"
pylint = "^2.15.4"
ipython = "^8.5.0"

　早速ですがpandasとPolarsそれぞれについて、今回の要件を実現するプログラムを記述します。

# aggregate_query.py
from typing import Union
import numpy as np
import pandas as pd
import polars as pl
def by_pandas(df: pd.DataFrame):
    """pandasを用いてアヤメの種類ごとの特徴量を抽出する"""
    features = (
        df
        .groupby('class')  # アヤメの種類で集約する
        # pandasではマルチインデックスがサポートされているのでこう書ける
        .agg({
            'sepal_length': [np.mean, np.std],
            'sepal_width': [np.mean, np.std],
            'petal_length': [np.mean, np.std],
            'petal_width': [np.mean, np.std],
        })
        # がく片の長さの平均値で表を並び替える（降順）
        .sort_values(by=[('sepal_length', 'mean')], ascending=False)
    )
    return features
def by_polars(df: Union[pl.DataFrame, pl.LazyFrame]):
    """polarsを用いてアヤメの種類ごとの特徴量を抽出する"""
    features = (
        df
        .groupby('class')  # アヤメの種類で集約する
        .agg(
            [
                # polarsはマルチインデックスを現在サポートしていない
                # リスト内包表記などを使うことで抽象化は可能
                pl.col('sepal_length').mean().alias('sepal_length_mean'),
                pl.col('sepal_length').std().alias('sepal_length_std'),
                pl.col('sepal_width').mean().alias('sepal_width_mean'),
                pl.col('sepal_width').std().alias('sepal_width_std'),
                pl.col('petal_length').mean().alias('petal_length_mean'),
                pl.col('petal_length').std().alias('petal_length_std'),
                pl.col('petal_width').mean().alias('petal_width_mean'),
                pl.col('petal_width').std().alias('petal_width_std'),
            ]
        )
        # がく片の長さの平均値で表を並び替える（降順）
        .sort('sepal_length_mean', reverse=True)
    )
    return features
if __name__ == '__main__':
    # pandasでの特徴量計算
    features_by_pandas = by_pandas(pd.read_csv('../data/irisx10000.csv'))
    print(features_by_pandas)
    # polarsでの特徴量計算（遅延評価しない）
    features_by_eager_polars = by_polars(pl.read_csv('../data/irisx10000.csv'))
    print(features_by_eager_polars)
    # polarsでの特徴量計算（遅延評価する）
    features_by_lazy_polars = by_polars(pl.scan_csv('../data/irisx10000.csv')).collect()
    print(features_by_lazy_polars)

　pandasもPolarsもDataFrameというクラスを持っており、これが2次元表を表すためのデータ構造になります。

データフレームを出力した結果です。単なる2次元表として取り扱うことができます。pandasの表示では左に連番が振られた列（インデックス列）がありますが、Polarsの出力にはありません

　by_pandas関数ではpandasのAPIを用いてクエリを、by_polars関数ではpolarsのAPIを用いてクエリを記述しています。提供されているAPIによる違いはありますが、どちらもclass（アヤメの種類）列でgroupby（集約）をして、集約関数としてmean（平均）とstd（標準偏差）を適用し、その出力列でソートをするという概略を素直に書くことができています。慣れる必要はありますが、SQLを書くことができる人であれば習得するのはそこまで難しくはないでしょう。

　PolarsにはpandasにはないLazyFrameというクラスが存在しており、この利用が推奨されています。これは計算の実行を極力遅延させ、collectメソッドを呼び出したときに計算を実行しDataFrameを返却するという機能を有しています。上記のプログラムにおいては、read_csvメソッドをscan_csvメソッドに置き換えるだけで、LazyFrameクラスを用いた遅延実行が可能になります。この遅延実行により、中間オブジェクトの生成を抑えるなど処理最適化を行える余地が大きくなると考えられます。

　上記の3条件（pandasのDataFrameを用いる、PolarsのDataFrameを用いる、PolarsのLazyFrameを用いる）についてPythonの実行時間計測ツールtimeitを用いて、簡易計測を行った結果が以下となります。

3条件についてそれぞれ時間計測を行った時の出力

　結果としてPolarsはpandasよりもおよそ6倍速く実行できました。また大きな差はないですがLazyFrameを用いた方がわずかに速いという結果が得られました。集約クエリは一度集約キーで分割した後はそれぞれ独立な表として扱えるので、並列処理可能な余地が大きいと考えられます。またPolarsのバックエンドはRustで実装されており、物理コアの性能が引き出しやすいと考えることができそうです。

時系列クエリを試してみる

　集約クエリは並列処理させやすいので、次は時系列に対するクエリを試してみたいところです。以前のRustパフォーマンス検証記事で作った時系列データを用いて再び移動平均を計算してみようと思います。

# time_series_query.py
from typing import Union
import pandas as pd
import polars as pl
def by_pandas(df: pd.DataFrame, window_size: int):
    """pandasを用いて移動平均を計算する"""
    features = df.rolling(window_size).mean()
    return features
def by_polars(df: Union[pl.DataFrame, pl.LazyFrame], window_size: int):
    """polarsを用いて移動平均を計算する"""
    features = df.select(
        [
            pl.col('value').rolling_mean(window_size),
        ]
    )
    return features
if __name__ == '__main__':
    # pandasでの特徴量計算
    series_by_pandas = by_pandas(pd.read_csv('../data/time_series.csv'), 50)
    print(series_by_pandas)
    # polarsでの特徴量計算(遅延評価しない)
    series_by_eager_polars = by_polars(pl.read_csv('../data/time_series.csv'), 50)
    print(series_by_eager_polars)
    # polarsでの特徴量計算(遅延評価する)
    series_by_lazy_polars = by_polars(pl.scan_csv('../data/time_series.csv'), 50).collect()
    print(series_by_lazy_polars)

　pandasで移動平均を計算する場合、文字通り1行で処理が記述できてしまいます。この手軽さはやはりpandasの魅力を非常によく表しています。一方でPolarsで移動平均を計算する場合、DataFrame、LazyFrame（表データ）自体には移動窓計算を行うAPIがなく、Series（列データ）にその機能があるので少し複雑な表現になります。

　集約クエリの時と同様、3条件（pandasをDataFrameを用いる、PolarsのDataFrameを用いる、PolarsのLazyFrameを用いる）について時間計測を行った結果が以下となります。

3条件についてそれぞれ時間計測を行った時の出力

　結果としてpolarsはpandasよりもおよそ8倍実行できました。また集計クエリの時と同じくLazyFrameを用いた方がわずかに速いという結果が得られました。前回の記事で作ったプログラムの中では590msが最高という結果であったので、それと比較してもおよそ7倍速いです。

Rust APIを試す

　PolarsではRust APIも提供されているとのことなので、時系列クエリをRustでも記述してみます。Rustおよび今回利用するライブラリのバージョンは以下のようになります。こちらはバージョン管理ツールCargoを利用して表示しています。

[package]
authors = ["Naoki Fujita"]
edition = "2021"
license = "MIT"
name = "web_engineer_in_rust"
rust-version = "1.64.0"
version = "0.1.0"
[dependencies]
anyhow = "1.0.65"
polars = {version = "0.24.3", features = ["lazy", "csv-file", "rolling_window"]}
serde_json = "1.0.68"
[[bin]]
name = "time_series_query"
path = "src/time_series_query.rs"

　Pythonの時系列クエリプログラムとほぼ等価なRustプログラムは以下のようになります。

// time_series_query.rs
use polars::prelude::{col, DataFrame, Duration, LazyCsvReader, LazyFrame, RollingOptions};
// プロジェクトディレクトリからの相対パスを絶対パスにするユーティリティー
fn get_csv_path(relative_path: &str) -> std::path::PathBuf {
    let project_path = env!("CARGO_MANIFEST_DIR");
    std::path::Path::new(project_path)
        .parent()
        .unwrap()
        .join(relative_path)
}
// polarsで移動平均を計算する
fn by_polars(df: LazyFrame, window_size: i64) -> anyhow::Result<DataFrame> {
    let duration = Duration::new(window_size);
    let rolling_options = RollingOptions {
        window_size: duration,
        min_periods: window_size as usize,
        ..RollingOptions::default()
    };
    let features = df
        .select([col("value").rolling_mean(rolling_options)])
        .collect()?;
    Ok(features)
}
fn main() -> anyhow::Result<()> {
    let csv_path = get_csv_path("data/time_series.csv");
    // csvを遅延読み込みする
    let df = LazyCsvReader::new(csv_path).has_header(true).finish()?;
    // polarsでの特徴量計算（遅延評価する）
    let features = by_polars(df, 50)?;
    println!("{:?}", features);
    Ok(())
}

　Python版とAPIの対比が取りやすいようにプログラムを構成しました。RustとPythonの大きな違いの一つとして、Pythonはキーワード引数という非常に利便性の高い関数呼び出し方法がありますが、Rustにはパフォーマンスを最大化するためにそのような機能が提供されていないことが挙げられます。結果として上記のLazyCsvReaderのように一つ一つのオプションを逐次的に選択するようなAPI（Builderパターンとも呼ばれる）が提供されていることが多く、そのような前提がAPI設計の差異を形成しています。

　そのような差異があることとAPIの型に合わせる必要性があることを除けば、Pythonとおおむね同じプログラム構成でRustからもPolarsを利用できます。一方でRust APIはドキュメントがPython APIほど充実してはおらず、基本的にはPythonからの利用を想定したライブラリであると思われます。

まとめ

　今回は超高速データ分析ライブラリPolarsを用いた集約クエリ、時系列移動平均計算クエリの記述方法を示しました。 さらにpandas APIとの差異や速度検証を行い、またRust APIを用いたクエリの記述方法も示しました。

　PolarsはベースがRustで設計されていることもあり、pandasと比べてAPIが「かたい」印象を受けますが、それにより卓越した高速化を実現しているため、より大きなデータセットをスムーズに分析したいというニーズに適したライブラリであると考えられます。

筆者紹介

藤田直己

1988年生まれ、大阪府枚方市出身

京都大学工学部電気電子工学科卒、同大学エネルギー科学研究科修了

応用情報技術者・ネットワークスペシャリスト・情報処理安全確保支援士試験合格者

YKK APにて超高層建築物の外装設計に従事し、型・モジュール設計・ウオーターフォールプロセスに精通する。その後ITエンジニアに転向。paizaにて、Ruby on RailsやReactを用いたWebサービスのスクラム開発に従事、現在に至る。

最も得意な言語はPython、最も影響を受けた言語はClojureであり、シンプルな関数型（的書き方ができる）言語を好む。関数型的記法を持ちながら、実行性能が高いRustに興味を持ち研さんを続けている。