2023年2月: Rust製高速データフレームライブラリ、Polarsを試す(kadowaki)

門脇@satoru_kadowakiです。 今月のPython Monthly Topicsでは、Rust製の高速データフレームライブラリ Polars について紹介します。

Polarsとは

Pythonでデータ分析に使用される主なライブラリに pandas があります。 Polarsはpandasと同様にデータフレームというデータ構造オブジェクトを提供するサードパーティライブラリです。 特にpandasを意識して作られており、メインページに「Lightning-fast DataFrame library for Rust and Python」とあるように、Rustによる高速処理を謳っています。

Polarsのリポジトリや関連ドキュメントは以下を参照してください。

Rust製でPythonバインディングを備えており、そのパフォーマンスについてはいくつかのベンチマークで高速なデータフレームライブラリの一つであることが示されています。

本記事では、Polarsとpandasのいくつかの機能を使用してその違いを見ていきます。

Polarsの特徴

まずはPolarsの特徴から見ていきます。下記は主にpandasを使用したことがある人を対象に挙げています。 以降では、実際にPolarsを使ったコードを示しながら、これらの特徴についても触れていきます。

  • Rust製で高速

  • pandasで使用されるメソッドと同じものが多く、pandas経験者にやさしい

  • 「Polars Expressions」という各種メソッドをつなぎ合わせて、データフレームの操作を行うことをコンセプトとしている

  • インデックスがない(pandasにはあるが、Polarsはインデックスがないことをメリットとしている)

  • 遅延評価(Lazy Evaluation)ができる

    • pandasは先行評価(Eager Evaluation)のみ

まずは使ってみよう

執筆時点での筆者が使用したPython、Polarsおよびpandasのバージョンは以下のとおりです。

  • Python 3.11.1

  • Polars 0.16.1

  • pandas 1.5.3

インストール

Polarsもpandasも pip コマンドで簡単にインストールできます。

$ pip install polars
$ pip install pandas

基本的な使い方

まずはデータフレームの作成から試してみます。 pandasでは以下のように行います。

>>> import pandas as pd
>>> data = {
...     "writer": ["kadowaki", "terada", "takanory", "ryu22e", "fukuda"],
...     "value": [1, 2, 3, 4, 5],
... }
>>> pd_df = pd.DataFrame(data)
>>> pd_df
     writer  value
0  kadowaki      1
1    terada      2
2  takanory      3
3    ryu22e      4
4    fukuda      5

Polarsでも以下のようにpandasと同様に行えます。 データフレームの表示はshapeとヘッダーに列の型を表示してくれています。

>>> import polars as pl
>>> pl_df = pl.DataFrame(data)  # dataはpandasで使用したもの
>>> pl_df
shape: (5, 2)
┌──────────┬───────┐
│ writer   ┆ value │
│ ---      ┆ ---   │
│ str      ┆ i64   │
╞══════════╪═══════╡
│ kadowaki ┆ 1     │
│ terada   ┆ 2     │
│ takanory ┆ 3     │
│ ryu22e   ┆ 4     │
│ fukuda   ┆ 5     │
└──────────┴───────┘

Polars Expressionsについて

データフレームが作成できたところで、Polarsの特徴である Polars Expressions について説明します。

Polars Expressionsとは、データフレームを操作するためのメソッド群です。 Polarsでは高速化のためにメソッドの引数に式を使用して処理することをコンセプトとしており、メソッドを連結して記述することで複雑な処理も高速に解決することを強みとしています。 [1]

具体的にどのようなことか、データフレームに列を追加する処理を例に説明します。

列の追加

列の追加は、pandasでは下記のように行えます。

>>> pd_df["tenx_value"] = pd_df["value"] * 10
>>> pd_df
     writer  value  tenx_value
0  kadowaki      1          10
1    terada      2          20
2  takanory      3          30
3    ryu22e      4          40
4    fukuda      5          50

Polarsで同様に行おうとすると、以下のようにエラーが発生します。

>>> pl_df["tenx_value"] = pl_df["value"] * 10
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/envs/py311_1/python3.11/site-packages/polars/internals/dataframe/frame.py", line 1573, in __setitem__
    raise TypeError(
TypeError: 'DataFrame' object does not support 'Series' assignment by index. Use 'DataFrame.with_columns'

同様に pl_df["value2"] = [11, 12, 13, 14, 15] のように、リストでSeries(シリーズ)のデータ構造を投入しようとすることもエラーになります。

Polarsで列の追加を行うには、 with_columns() メソッドを使用します。 [2] 引数として polars.col() メソッドで列名を指定し、列の別名割り当てに alias() メソッドを使用します。

具体的なコードは以下のようになります。

>>> pl_df = pl_df.with_columns([(pl.col("value") * 10).alias("tenx_value")])
>>> pl_df
shape: (5, 3)
┌──────────┬───────┬────────────┐
│ writer   ┆ value ┆ tenx_value │
│ ---      ┆ ---   ┆ ---        │
│ str      ┆ i64   ┆ i64        │
╞══════════╪═══════╪════════════╡
│ kadowaki ┆ 1     ┆ 10         │
│ terada   ┆ 2     ┆ 20         │
│ takanory ┆ 3     ┆ 30         │
│ ryu22e   ┆ 4     ┆ 40         │
│ fukuda   ┆ 5     ┆ 50         │
└──────────┴───────┴────────────┘

pandasでよく使用される機能との比較

ここからは、より実践的なデータを使用しながらpandasとの違いをみていきます。 サンプルスクリプトでは 気象庁ホームページ で公開されている 「最新の気象データ」24時間降水量一覧表ページ に公開されている CSVデータ をダウンロードし、使用しています。 (本データはおよそ10分毎に更新されています。)

CSVファイルの読み込み

最初にCSVファイルを読み込んでみます。 import文のあとに表示オプションを指定していますが、具体的な説明は割愛します。

オプションパラメーターについては、 それぞれのAPIリファレンス(polars.read_csv, pandas.read_csv)に記載がありますので、興味のある方は読んでみてください。

  • pandasの例

pd_readcsv.py
import pandas as pd

pd.options.display.unicode.east_asian_width = True  # print()で等幅を指定
pd.options.display.max_columns = 8  # 表示カラム数の指定
pd.options.display.width = 200  # 表示幅の指定

pd_df = pd.read_csv("./preall00_rct.csv", encoding="shiftjis")  # CSVのエンコードを指定
print(pd_df.head(5))
  • Polarsの例

pl_readcsv.py
import polars as pl

pl_df = pl.read_csv("./preall00_rct.csv", encoding="shiftjis")
print(pl_df.head(5))

細かいオプションを使用せずシンプルに読み込むだけであれば、2つのコードは全く同じと言えます。 それぞれの実行結果は以下のとおりです。

  • pandas

$ python3.11 pd_readcsv.py
   観測所番号         都道府県                        地点  国際地点番号  ...  72時間降水量 今日の最大値(mm)  72時間降水量 今日の最大値の品質情報  日降水量 今日の値(mm)  日降水量 今日の値の品質情報
0       11001  北海道 宗谷地方      宗谷岬(ソウヤミサキ)           NaN  ...                            8.0                                    4                    0.0                            4
1       11016  北海道 宗谷地方          稚内(ワッカナイ)       47401.0  ...                           19.0                                    4                    0.5                            4
2       11046  北海道 宗谷地方              礼文(レブン)           NaN  ...                           18.0                                    4                    2.0                            4
3       11061  北海道 宗谷地方            声問(コエトイ)           NaN  ...                           10.5                                    4                    0.5                            4
4       11076  北海道 宗谷地方  浜鬼志別(ハマオニシベツ)           NaN  ...                            5.5                                    4                    0.0                            4

[5 rows x 55 columns]
  • Polars

$ python3.11 pl_readcsv.py
shape: (5, 55)
┌────────────┬─────────────────┬────────────────────────────┬──────────────┬─────┬───────────────────────────────┬─────────────────────────────────────┬───────────────────────┬─────────────────────────────┐
│ 観測所番号  都道府県         地点                        国際地点番号  ...  72時間降水量 今日の最大値(mm)  72時間降水量 今日の最大値の品質情報  日降水量 今日の値(mm)  日降水量 今日の値の品質情報 │
│ ---         ---              ---                         ---                ---                            ---                                  ---                    ---                         │
│ i64         str              str                         i64                f64                            i64                                  f64                    i64                         │
╞════════════╪═════════════════╪════════════════════════════╪══════════════╪═════╪═══════════════════════════════╪═════════════════════════════════════╪═══════════════════════╪═════════════════════════════╡
│ 11001       北海道 宗谷地方  宗谷岬(ソウヤミサキ)      null          ...  8.0                            4                                    0.0                    4                           │
│ 11016       北海道 宗谷地方  稚内(ワッカナイ)          47401         ...  19.0                           4                                    0.5                    4                           │
│ 11046       北海道 宗谷地方  礼文(レブン)              null          ...  18.0                           4                                    2.0                    4                           │
│ 11061       北海道 宗谷地方  声問(コエトイ)            null          ...  10.5                           4                                    0.5                    4                           │
│ 11076       北海道 宗谷地方  浜鬼志別(ハマオニシベツ)  null          ...  5.5                            4                                    0.0                    4                           │
└────────────┴─────────────────┴────────────────────────────┴──────────────┴─────┴───────────────────────────────┴─────────────────────────────────────┴───────────────────────┴─────────────────────────────┘

なお、どちらの出力結果もターミナルに表示された内容をコピー&ペーストしています。全角文字が含まれるため、表示上の位置ずれが発生していますが、実際のターミナルでは等幅フォントを使用することで位置ずれなく表示されるようです。 以下に出力結果の参考画像を掲載します。

  • pandasの例 pandasの例

  • Polarsの例 Polarsの例

CSVの読み込み速度、メモリ使用量を比較

「PolarsはRust製だから速い」と言われるものの、実際どの程度違うのかCSVファイルの読み込み速度とメモリ使用量を簡単に比較してみました。

  • 計測方法

  • テストに使用するデータ

    • 前述のCSVファイルの読み込みで使用した CSV を繰り返し連結して作成

      • 合計: 2,631,680行

      • ファイルサイズ: 約430MB

計測に使用したスクリプトは以下のとおりです。

  • pandasの場合

pd_compare_readcsv.py
import pandas as pd
from memory_profiler import profile


@profile
def main():
    pd_df = pd.read_csv("./test.csv")


if __name__ == "__main__":
    main()
  • Polarsの場合

pl_compare_readcsv.py
import polars as pl
from memory_profiler import profile


@profile
def main():
    pl_df = pl.read_csv("./test.csv")


if __name__ == "__main__":
    main()

上記のコードをファイルに保存し、timeコマンドを使用して以下のように実行します。

$ time python3.11 pd_compare_readcsv.py
$ time python3.11 pl_compare_readcsv.py

結果は以下のようにPolarsの方がメモリ使用量も、処理時間も圧倒的によい結果となりました。 シンプルに速いというだけで嬉しくなりますね。

項目

結果: pandas

結果: Polars

メモリ使用量(MB)

2163.17

1371.65

処理にかかった時間(real)

27.8s

5.4s

ユーザーCPU時間(user)

17.1s

10.5s

システムCPU時間(sys)

4.2s

3.1s

行や列の選択

行や列の選択について説明する前に、pandasとPolarsの大きな違いの一つである「インデックスの有無」について説明します。 pandasでは角括弧にインデックスを指定して行や列の選択を行うことがよくありますが、Polarsではそもそもインデックス自体が存在しません。 pandas同様に角括弧を使用した選択も行うことができますが、この方法はアンチパターンとされており、将来使用できなくなる可能性があるとされています。

また、インデックスを使用しないことのメリットとして以下のようなことがあります。

  • インデックスを使用しないことで遅延評価ができる

    • インデックスは先行評価しかできない

    • (遅延評価と先行評価については後述します)

  • 複数列の操作をはじめ、多くの処理の並列化を実現している

    • 角括弧を使用したインデックスの操作はシングルスレッドしかできない

下記のユーザーガイドでも、インデックスが有効なケースを除き、後述するメソッド群を使用することを推奨していますので読んでみてください。

前置きが長くなりましたが、具体的な方法を先述のpandasとPolarsのコードに追加してみていきます。

pandasではインデックスn番目からm番目を指定したスライスは以下のように行います。

print(pd_df.loc[375:379, ["都道府県", "地点", "24時間降水量 現在値(mm)"]])
  • 出力結果

    都道府県                  地点  24時間降水量 現在値(mm)
375   山形県      櫛引(クシビキ)                     25.0
376   山形県      肘折(ヒジオリ)                     19.5
377   山形県  尾花沢(オバナザワ)                     12.0
378   山形県  鼠ケ関(ネズガセキ)                     10.5
379   山形県      荒沢(アラサワ)                     23.5

Polarsでは、行方向と列方向にそれぞれのメソッドを使用して選択します。 以下のコードは、列を選択するために select() メソッドを使用し、指定した行位置から末尾の行を取得するために head() メソッドと tail() メソッドを組み合わせて使用しています。

print(pl_df.select(pl.col(["都道府県", "地点", "24時間降水量 現在値(mm)"]).head(380).tail(5)))
  • 出力結果

shape: (5, 3)
┌────────────┬──────────────────────┬─────────────────────────┐
│ 都道府県    地点                  24時間降水量 現在値(mm) │
│ ---         ---                   ---                     │
│ str         str                   f64                     │
╞════════════╪══════════════════════╪═════════════════════════╡
│ 山形県      櫛引(クシビキ)      25.0                    │
│ 山形県      肘折(ヒジオリ)      19.5                    │
│ 山形県      尾花沢(オバナザワ)  12.0                    │
│ 山形県      鼠ケ関(ネズガセキ)  10.5                    │
│ 山形県      荒沢(アラサワ)      23.5                    │
└────────────┴──────────────────────┴─────────────────────────┘

また、pandasでは条件式を使用してデータを抽出することがあります。 以下は列「都道府県」が「山形県」である行の特定の列「都道府県、地点、24時間降水量 現在値(mm)」を抽出している例です。

print(pd_df.loc[(pd_df["都道府県"] == "山形県"), ["都道府県", "地点", "24時間降水量 現在値(mm)"]])
  • 出力結果

    都道府県                            地点  24時間降水量 現在値(mm)
364   山形県                飛島(トビシマ)                      7.0
365   山形県                  酒田(サカタ)                     14.0
...(省略)
390   山形県                高峰(タカミネ)                     24.0
391   山形県                米沢(ヨネザワ)                     16.5

Polarsで同様に行うには filter() メソッドを使用します。

print(pl_df.select(pl.col(["都道府県", "地点", "24時間降水量 現在値(mm)"])).filter(pl.col("都道府県") == "山形県"))
  • 出力結果

shape: (28, 3)
┌────────────┬────────────────────────────┬─────────────────────────┐
│ 都道府県    地点                        24時間降水量 現在値(mm) │
│ ---         ---                         ---                     │
│ str         str                         f64                     │
╞════════════╪════════════════════════════╪═════════════════════════╡
│ 山形県      飛島(トビシマ)            7.0                     │
│ 山形県      差首鍋(サスナベ)          28.5                    │
│ ...         ...                         ...                     │
│ 山形県      高峰(タカミネ)            24.0                    │
│ 山形県      米沢(ヨネザワ)            16.5                    │
└────────────┴────────────────────────────┴─────────────────────────┘

データフレームの結合、マージ

データフレームの結合やマージで使用されるメソッドは以下の表の通りです。 Polarsでは merge() メソッドがなくマージと結合には join() メソッドが使用されます。

処理

pandasで使用されるメソッド

Polars

連結

pandas.concat()

polars.concat()

マージ

pandas.merge()

polars.DataFrame.join()

結合

pandas.DataFrame.join()

polars.DataFrame.join()

基本的な使い方が似ているため、ここではPolarsで行う方法についてのみ紹介します。

連結

pl_concat.py
import polars as pl

df1 = pl.DataFrame(
    {
        "writer": ["kadowaki", "terada", "takanory"],
        "value": [1, 2, 3],
    }
)
df2 = pl.DataFrame(
    {
        "writer": ["ryu22e", "fukuda"],
        "value": [4, 5],
    }
)
print(pl.concat([df1, df2], how="vertical"))

上記のコードを実行すると以下の結果になります。

$ python3.11 pl_concat.py
shape: (5, 2)
┌──────────┬───────┐
│ writer    value │
│ ---       ---   │
│ str       i64   │
╞══════════╪═══════╡
│ kadowaki  1     │
│ terada    2     │
│ takanory  3     │
│ ryu22e    4     │
│ fukuda    5     │
└──────────┴───────┘

マージ、結合

Polarsでデータフレームのマージや結合を行うには先述のとおり polars.DataFrame.join() を使用します。 pandasの join() メソッドはインデックスをもとにして3つ以上のデータフレームを連結できますが、Polarsではデータフレームの連結は2つまでという違いがあります。どちらかというと、pandasの merge() メソッドと同等と考えるのがよさそうです。

pl_merge.py
import polars as pl

df1 = pl.DataFrame(
    {
        "writer": ["kadowaki", "terada", "takanory", "ryu22e"],
        "value": [1, 2, 3, 4],
    }
)
df2 = pl.DataFrame(
    {
        "writer": ["ryu22e", "fukuda", "kadowaki"],
        "value": [5, 6, 7],
    }
)
print(df1.join(df2, on="writer", how="inner"))

結果は以下のようになります。

$ python3.11 pl_merge.py
shape: (2, 3)
┌──────────┬───────┬─────────────┐
│ writer   ┆ value ┆ value_right │
│ ---      ┆ ---   ┆ ---         │
│ str      ┆ i64   ┆ i64         │
╞══════════╪═══════╪═════════════╡
│ kadowaki ┆ 1     ┆ 7           │
│ ryu22e   ┆ 4     ┆ 5           │
└──────────┴───────┴─────────────┘

また、これらの処理速度についてもpandasと比較してみたところ以下のような結果となりました。 連結やマージでも処理時間の差が見られます。

pandasコードサンプル

pandas処理時間

Polars処理時間

pd_concat.py

real: 0.658s, user: 0.699s, sys: 0.241s

real: 0.180s, user: 0.136s, sys: 0.032s

pd_merge.py

real: 0.700s, user: 0.726s, sys: 0.285s

real: 0.210s, user: 0.123s, sys: 0.064s

列に対する処理

pandasで行や列に対して処理を行いたい場合に apply() メソッドが使用されます。Polarsにも同じ名前のメソッドがあり、pandasと同じように使用できます。

pl_apply.py
import polars as pl

df = pl.DataFrame({"value": [1, 2, 3, 4, 5]})

# 偶数か奇数かを判定
def even_odd(x):
    if x % 2 == 0:
        return "Even"
    else:
        return "Odd"


# 列valueに演算を行い、tenx_valueカラムを追加
df = df.with_columns(pl.col("value").apply(lambda x: even_odd(x)).alias("even_odd"))
print(df)

実行結果は以下になります。

$ python3.11 pl_apply.py
shape: (5, 2)
┌───────┬──────────┐
│ value  even_odd │
│ ---    ---      │
│ i64    str      │
╞═══════╪══════════╡
│ 1      Odd      │
│ 2      Even     │
│ 3      Odd      │
│ 4      Even     │
│ 5      Odd      │
└───────┴──────────┘

日付やdict型への変換処理

文字列を日付型に変更する場合は、pandasでは pandas.to_datetime() メソッドを使用します。 Polarsでは下記のように .str.strptime(pl.Datetime, fmt="%Y-%m-%d")) のようにメソッドチェーンを使用して行います。 pl.Datetime の部分を pl.Date のようにすればdate型に変換されます。

pl_todatetime.py
import polars as pl

df = pl.DataFrame({"someday": ["1956-01-31", "1991-02-20", "2015-05-16"]})
df = df.with_columns(pl.col("someday").str.strptime(pl.Datetime, fmt="%Y-%m-%d"))
print(df)
  • 実行結果

$ python3.11 pl_todatetime.py
shape: (3, 1)
┌─────────────────────┐
│ someday             │
│ ---                 │
│ datetime[μs]        │
╞═════════════════════╡
│ 1956-01-31 00:00:00 │
│ 1991-02-20 00:00:00 │
│ 2015-05-16 00:00:00 │
└─────────────────────┘

また、pandasでデータフレームをdict型やJSONに変換する場合は pandas.DataFrame.to_dict() あるいは to_json() メソッドを使用します。 Polarsではdict型の場合は polars.DataFrame.to_dict() メソッドを使用し、JSONに変換する場合は polars.DataFrame.write_json() メソッドを使用します。

ここでは、dict型に変換する方法を紹介します。 as_series オプションにTrue(デフォルト値)を指定すると、値が polars.internals.series.series.Series クラスのインスタンスとして出力されるため、シリーズで出力したくない場合はこのオプションをFalseにします。

pl_todict.py
import polars as pl

df = pl.DataFrame(
    {
        "writer": ["kadowaki", "terada", "takanory"],
        "value": [1, 2, 3],
    }
)
print(df.to_dict(as_series=False))
  • 実行結果

$ python3.11 pl_todict.py
{'writer': ['kadowaki', 'terada', 'takanory'], 'value': [1, 2, 3]}

遅延評価

遅延評価(Lazy Evaluation)とは、ある式をすぐに評価せず必要なときにだけ評価することです。 宣言した時点では何も処理が行われないことから、Pythonのジェネレーターの考え方に近いと思いますが、Polarsでは遅延評価を行うためのメソッドやオプションが用意されています。

pandasでは式を即時に実行する先行評価(Eager Evaluation)のみが利用できます。 そのため、処理するデータサイズに合わせてメモリが使用されますが、Polarsでは polars.LazyFrame() メソッドにより遅延評価用のデータフレームを使用することでメモリを効率的に使用できます。

具体的にみていきましょう。CSVを読み込む read_csv() メソッドについては先述の通りですが、これをLazyFrameにしてみます。方法は簡単で scan_csv() メソッドを使用するだけです。

>>> df = pl.scan_csv("./lazytest.csv")
>>> df
<polars.LazyFrame object at 0x7FA564C4D090>

また、LazyDataFrameに対して filter() メソッドなどを実行することもできますが、メソッドはすぐには実行されません。 実行するには、データを取得するための fetch()collect() メソッドを使用する必要があります。 これらのメソッドが実行されたタイミングで初めて式が評価されます。

>>> df = df.select(pl.col(["都道府県", "地点"])).filter(pl.col("都道府県") == "山形県")
>>> df
<polars.LazyFrame object at 0x7F44D20C1610>
>>> df.fetch()  # デフォルトでは先頭500行が取得される
shape: (56, 2)
┌────────────┬────────────────────────────┐
│ 都道府県   ┆ 地点                       │
│ ---        ┆ ---                        │
│ str        ┆ str                        │
╞════════════╪════════════════════════════╡
│ 山形県     ┆ 飛島(トビシマ)           │
│ 山形県     ┆ 酒田(サカタ)             │
│ ...        ┆ ...                        │
│ 山形県     ┆ 高峰(タカミネ)           │
│ 山形県     ┆ 米沢(ヨネザワ)           │
└────────────┴────────────────────────────┘
>>> df.collect()
shape: (56, 2)
...(省略)

読み込み済みのデータフレームに対しても遅延評価を行うことができます。 こちらもやり方は簡単で lazy() メソッドを使用してメソッドを繋ぐだけです。

# dfに対して.lazy().select()...のようにメソッドチェーンを行う
>>> df = df.lazy().select(pl.col(["都道府県", "地点"])).filter(pl.col("都道府県") == "山形県")
>>> df
<polars.LazyFrame object at 0x7F44D20DE450>
>>> df.collect()
shape: (56, 2)
...(省略)

遅延評価がこんなに簡単に使えるのは便利ですね。 ファイル読み込みやクエリを先に宣言しておき、必要なタイミングで値を取得できるので処理の最適化が図れそうです。

まとめ

いかがでしたか? Polarsはpandasを意識して作成されていることから、シンプルなものならパフォーマンス向上を図るためだけにpandasからの乗り換えを考えるのも悪くなさそうです。 また、遅延評価を使用することでメモリを節約しながら大きいデータを処理したいケースなどにおいても試してみる価値がありそうです。 しかし、Polarsは最初のリリースから1年未満という新しいライブラリで、pandasの機能をすべてカバーしているわけではありません。乗り換えには十分な評価を行うことをおすすめします。 Polarsの今後にも期待しながら使いどころを探ってみてください!