🐻‍❄️

Pandasを速くしたい? Polarsを知ろう・触ろう!

2023/07/14に公開

Pandasオルタナティブの一角で、Pandasよりも高速と噂されているPolarsについてまとめました。
https://www.pola.rs/

主に以下の公式ドキュメントを参考にしてます。

なお、本記事では性能検証やPandas APIとの比較などについては行いません。

特徴

Polarsはテーブルデータを取り扱うためのライブラリです。

他のOSSではPandasやdaskなどが類似ライブラリとなりますが、Introductionではそれらと比較したときのPolarsの特徴として、以下が挙げられています。

なぜ高速でメモリ効率も良いのか?

https://www.pola.rs/posts/i-wrote-one-of-the-fastest-dataframe-libraries/ にそのコンセプトがまとめられています。

Arrow

Arrowは、言語やランタイムに依存しないカラム型メモリフォーマットで、巨大なデータを相手にする分析等のタスクを想定しています。

  • データは連続するメモリバッファに格納される
    • キャッシュミス率が下がる
    • SIMDで効率化させやすい
    • 文字列などの可変長のデータはオフセット配列を別に持つ
  • テーブルのカラムのメモリバッファのオフセットとサイズをヘッダとして保持している
    • RAMに乗らない巨大なデータセットの一部を読み込む、などが可能
    • ランダムなリードアクセスが得意なSSDとの親和性が高い
  • 参照カウンタを持ち、イミュータブル
    • コピーが発生することがほとんどない
    • スライス操作も、参照カウンタのインクリメントとオフセットの変更のみ
  • データ要素ごとに欠損値かどうかをメタデータとしてビットで持つ (null-bit buffer)
    • 欠損値の表現がデータ型に依らない
    • ハードウェア演算子やSIMDによって欠損値を含む処理が高速化できる
    • nullカウントもメタデータに持つため、0ならnullチェックの分岐処理は不要

遅延評価

Polarsには、クエリの遅延評価がサポートされています。
コードを実行してもすぐに実行せず、クエリプランに追加していき、最終的にデータを取得・加工しようとしたときにクエリプランを最適化してから実行されます。
この遅延評価は、後述するExprやコンテキストによって支えられています。

TPC-Hでのベンチマーク

https://www.pola.rs/benchmarks.html ではデータベースに用いられるTPC-Hでベンチマークした結果が公開されてますが、クエリの計算時間もメモリ消費量もPandasと比較にならないくらいに効率的です。

クエリの実行時間

https://www.pola.rs/benchmarks.html から転載 (2023-07-07時点)

クエリのメモリ消費

https://www.pola.rs/benchmarks.html から転載 (2023-07-07時点)

Pandasとの比較

https://pola-rs.github.io/polars-book/user-guide/migration/pandas/

  • インデックスを持たない (もちろんマルチインデックスも無い)
    • Polarsではテーブル内の位置でインデックス付けされる
  • メモリ上のデータはNumpy配列ではなく、Arrow配列
  • 多くの処理が並列実行可能
    • Pandasは基本的にすべて直列処理
  • 遅延評価をサポート
    • Pandasは常に即時評価

動かしながら手触りを確かめる

前提となる環境

Polars 0.18.4を使います。

$ python --version               
Python 3.11.3

$ cat requirements.txt                             
polars==0.18.4
numpy==1.25.0

$ pip install -r requirements.txt

polarsはplでエイリアスするのが一般的です。

import polars as pl
import numpy as np

データ構造

Pandasなどと同様、Seriesが1次元、DataFrameが2次元 (0個以上のSeriesの集合) のデータ構造となってます。

s1 = pl.Series("a", [1, 2, 3])
print(s1)
# shape: (3,)
# Series: 'a' [i64]
# [
# 	1
# 	2
# 	3
# ]

s2 = pl.Series("b", ["a", "b", "c"])
df = pl.DataFrame([s1, s2])
print(df)
# shape: (3, 2)
# ┌─────┬─────┐
# │ a   ┆ b   │
# │ --- ┆ --- │
# │ i64 ┆ str │
# ╞═════╪═════╡
# │ 1   ┆ a   │
# │ 2   ┆ b   │
# │ 3   ┆ c   │
# └─────┴─────┘

基本的にはArrowを踏襲してますが、CategoricalとObjectがPolars特有の実装になっています。

  • 整数型: Int8, Int16, Int32, Int64, UInt8, Uint16, UInt32, UInt64
  • 浮動小数点型: Float32, Float64
  • 日時型: Date, DateTime, Duration, Time
  • ブール型: Boolean (ビット配列にまとめられる)
  • List
  • Struct (1カラムに複数のSeriesを含む、後述)
  • Binary
  • Utf8 (ArrowのLargeUtf8に対応)
  • Categorical (文字列カラムを効率的に扱うためにエンコーディングしたもの)
  • Object (任意のPythonオブジェクトをラップする)
    • サポートも限定的で、進んで利用すべきではない

ファイルからの読み込み

CSV、JSON、Parquetなどの多くのファイルをサポートをしています。
ファイルタイプごとにread_*とscan_*の2種類の関数が提供されています。例えばCSVファイルの場合はread_csv()scan_csv()です。
scan_*を使うと遅延ロードになります。クエリを解析して、ロードする行・列を予め削減しやすくなります。

eager_df = pl.read_csv("data/titanic.csv")
# shape: (9, 13)
# ┌────────────┬─────────────┬──────────┬──────────┬───┬───────────┬───────────┬───────┬──────────┐
# │ describe   ┆ PassengerId ┆ Survived ┆ Pclass   ┆ … ┆ Ticket    ┆ Fare      ┆ Cabin ┆ Embarked │
# │ ---        ┆ ---         ┆ ---      ┆ ---      ┆   ┆ ---       ┆ ---       ┆ ---   ┆ ---      │
# │ str        ┆ f64         ┆ f64      ┆ f64      ┆   ┆ str       ┆ f64       ┆ str   ┆ str      │
# ╞════════════╪═════════════╪══════════╪══════════╪═══╪═══════════╪═══════════╪═══════╪══════════╡
# │ count      ┆ 891.0       ┆ 891.0    ┆ 891.0    ┆ … ┆ 891       ┆ 891.0     ┆ 891   ┆ 891      │
# │ null_count ┆ 0.0         ┆ 0.0      ┆ 0.0      ┆ … ┆ 0         ┆ 0.0       ┆ 687   ┆ 2        │
# │ mean       ┆ 446.0       ┆ 0.383838 ┆ 2.308642 ┆ … ┆ null      ┆ 32.204208 ┆ null  ┆ null     │
# │ std        ┆ 257.353842  ┆ 0.486592 ┆ 0.836071 ┆ … ┆ null      ┆ 49.693429 ┆ null  ┆ null     │
# │ min        ┆ 1.0         ┆ 0.0      ┆ 1.0      ┆ … ┆ 110152    ┆ 0.0       ┆ A10   ┆ C        │
# │ max        ┆ 891.0       ┆ 1.0      ┆ 3.0      ┆ … ┆ WE/P 5735 ┆ 512.3292  ┆ T     ┆ S        │
# │ median     ┆ 446.0       ┆ 0.0      ┆ 3.0      ┆ … ┆ null      ┆ 14.4542   ┆ null  ┆ null     │
# │ 25%        ┆ 223.0       ┆ 0.0      ┆ 2.0      ┆ … ┆ null      ┆ 7.8958    ┆ null  ┆ null     │
# │ 75%        ┆ 669.0       ┆ 1.0      ┆ 3.0      ┆ … ┆ null      ┆ 31.0      ┆ null  ┆ null     │
# └────────────┴─────────────┴──────────┴──────────┴───┴───────────┴───────────┴───────┴──────────┘

lazy_df = pl.scan_csv("data/titanic.csv")
print(lazy_df)
# naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
#   CSV SCAN data/titanic.csv
#   PROJECT */12 COLUMNS

遅延させる場合、collect()して初めてDataFrameとしてメモリにロードされます。
(データのロードが必要なdescribe()などは、DataFrameでしかサポートされません)

print(type(lazy_df))
# <class 'polars.lazyframe.frame.LazyFrame'>

print(type(lazy_df.collect()))
# <class 'polars.dataframe.frame.DataFrame'>

引数dtypesで、予め型を指定することもできます。
PassengerIdをUInt16、SurvivedをCategoricalにしています。

df = pl.read_csv(
    "./data/titanic.csv",
    dtypes={
        "PassengerId": pl.UInt16,
        "Survived": pl.Categorical,
    },
)
df.schema
# {'PassengerId': UInt16,
#  'Survived': Categorical,
#  'Pclass': Int64,
#  'Name': Utf8,
#  'Sex': Utf8,
#  'Age': Float64,
#  'SibSp': Int64,
#  'Parch': Int64,
#  'Ticket': Utf8,
#  'Fare': Float64,
#  'Cabin': Utf8,
#  'Embarked': Utf8}

DataFrameを確認する

DataFrameのサイズはshape、スキーマはschemaで確認できます。
Pandas同様、descirbe()で各カラムのnullカウントや統計値などのサマリが表示されます。

titanic_df = pl.read_csv("data/titanic.csv")

titanic_df.shape
# (891, 12)

titanic_df.schema
# {'PassengerId': Int64,
#  'Survived': Int64,
#  'Pclass': Int64,
#  'Name': Utf8,
#  'Sex': Utf8,
#  'Age': Float64,
#  'SibSp': Int64,
#  'Parch': Int64,
#  'Ticket': Utf8,
#  'Fare': Float64,
#  'Cabin': Utf8,
#  'Embarked': Utf8}

titanic_df.describe()
# shape: (9, 13)
# ┌────────────┬─────────────┬──────────┬──────────┬───┬───────────┬───────────┬───────┬──────────┐
# │ describe   ┆ PassengerI… ┆ Survived ┆ Pclass   ┆ … ┆ Ticket    ┆ Fare      ┆ Cabin ┆ Embarked │
# │ ---        ┆ ---         ┆ ---      ┆ ---      ┆   ┆ ---       ┆ ---       ┆ ---   ┆ ---      │
# │ str        ┆ f64         ┆ f64      ┆ f64      ┆   ┆ str       ┆ f64       ┆ str   ┆ str      │
# ╞════════════╪═════════════╪══════════╪══════════╪═══╪═══════════╪═══════════╪═══════╪══════════╡
# │ count      ┆ 891.0       ┆ 891.0    ┆ 891.0    ┆ … ┆ 891       ┆ 891.0     ┆ 891   ┆ 891      │
# │ null_count ┆ 0.0         ┆ 0.0      ┆ 0.0      ┆ … ┆ 0         ┆ 0.0       ┆ 687   ┆ 2        │
# │ mean       ┆ 446.0       ┆ 0.383838 ┆ 2.308642 ┆ … ┆ null      ┆ 32.204208 ┆ null  ┆ null     │
# │ std        ┆ 257.353842  ┆ 0.486592 ┆ 0.836071 ┆ … ┆ null      ┆ 49.693429 ┆ null  ┆ null     │
# │ min        ┆ 1.0         ┆ 0.0      ┆ 1.0      ┆ … ┆ 110152    ┆ 0.0       ┆ A10   ┆ C        │
# │ max        ┆ 891.0       ┆ 1.0      ┆ 3.0      ┆ … ┆ WE/P 5735 ┆ 512.3292  ┆ T     ┆ S        │
# │ median     ┆ 446.0       ┆ 0.0      ┆ 3.0      ┆ … ┆ null      ┆ 14.4542   ┆ null  ┆ null     │
# │ 25%        ┆ 223.0       ┆ 0.0      ┆ 2.0      ┆ … ┆ null      ┆ 7.8958    ┆ null  ┆ null     │
# │ 75%        ┆ 669.0       ┆ 1.0      ┆ 3.0      ┆ … ┆ null      ┆ 31.0      ┆ null  ┆ null     │
# └────────────┴─────────────┴──────────┴──────────┴───┴───────────┴───────────┴───────┴──────────┘

ファイルへの書き出し

ファイルへの書き出しは、write_X()で行います。CSVであれば、write_csv()です。
Pandasと違って、to_X()ではない点に注意です。to_X()もありますが、こちらは他の型への変換になっています。(後述)

titanic_df.write_csv("titanic.csv")

Polarsの設定

Polarsの設定はConfigのset_X()で行います。

一時的に変更したい場合、withコンテキスト内やデコレータ内で書き換えます。
インタプリタ内で永続的に書き換えたい場合は、クラスメソッドを呼び出します。restore_defaults()でデフォルトに戻せます。

# NOTE: デフォルト表示
print(titanic_df.select([pl.col("Name")]))
# ┌───────────────────────────────────┐
# │ Name                              │
# │ ---                               │
# │ str                               │
# ╞═══════════════════════════════════╡
# │ Braund, Mr. Owen Harris           │
# │ Cumings, Mrs. John Bradley (Flor… │
# │ Heikkinen, Miss. Laina            │
# │ Futrelle, Mrs. Jacques Heath (Li… │
# │ …                                 │
# │ Graham, Miss. Margaret Edith      │
# │ Johnston, Miss. Catherine Helen … │
# │ Behr, Mr. Karl Howell             │
# │ Dooley, Mr. Patrick               │
# └───────────────────────────────────┘

with pl.Config() as cfg:
    # NOTE: with句内で文字列の表示長と表示行数を変更
    cfg.set_fmt_str_lengths(100)
    cfg.set_tbl_rows(4)
    print(titanic_df.select([pl.col("Name")]))
# ┌─────────────────────────────────────────────────────┐
# │ Name                                                │
# │ ---                                                 │
# │ str                                                 │
# ╞═════════════════════════════════════════════════════╡
# │ Braund, Mr. Owen Harris                             │
# │ Cumings, Mrs. John Bradley (Florence Briggs Thayer) │
# │ …                                                   │
# │ Behr, Mr. Karl Howell                               │
# │ Dooley, Mr. Patrick                                 │
# └─────────────────────────────────────────────────────┘

pl.Config.set_fmt_str_lengths(10)
print(titanic_df.select([pl.col("Name")]))
# ┌─────────────┐
# │ Name        │
# │ ---         │
# │ str         │
# ╞═════════════╡
# │ Braund, Mr… │
# │ Cumings, M… │
# │ Heikkinen,… │
# │ Futrelle, … │
# │ …           │
# │ Graham, Mi… │
# │ Johnston, … │
# │ Behr, Mr. … │
# │ Dooley, Mr… │
# └─────────────┘

# NOTE: デフォルトに戻す
pl.Config.restore_defaults()

DataFrameの操作

詳細に入る前に、Polarsでの集計コードをざっと見てみましょう。
PandasやSQLに馴染みがある方なら、パッとみただけで何を行なっているのかは、おおむね想像できるかと思います。
filterで行選択、selectで列選択、groupbyとaggで集約、sortで並び替えを行っています。

pl.read_csv("./data/titanic.csv").filter(
    (pl.col("Embarked") == "S") & (pl.col("Age").is_not_null())
).select([
    pl.col("Survived"),
    pl.col("Pclass"),
    (pl.col("Age") > 18).alias("is_adult"),
]).groupby(["Pclass", "is_adult"]).agg(
    pl.col("*").count().alias("count"),
    pl.col("Survived").mean().alias("survival_rate")
).sort(
    "Pclass", "is_adult", descending=[False, True]
)
# shape: (6, 4)
# ┌────────┬──────────┬───────┬───────────────┐
# │ Pclass ┆ is_adult ┆ count ┆ survival_rate │
# │ ---    ┆ ---      ┆ ---   ┆ ---           │
# │ i64    ┆ bool     ┆ u32   ┆ f64           │
# ╞════════╪══════════╪═══════╪═══════════════╡
# │ 1      ┆ true     ┆ 98    ┆ 0.581633      │
# │ 1      ┆ false    ┆ 10    ┆ 0.9           │
# │ 2      ┆ true     ┆ 131   ┆ 0.419847      │
# │ 2      ┆ false    ┆ 25    ┆ 0.76          │
# │ 3      ┆ true     ┆ 222   ┆ 0.189189      │
# │ 3      ┆ false    ┆ 68    ┆ 0.279412      │
# └────────┴──────────┴───────┴───────────────┘

Exprとコンテキスト

DataFrameの操作でキモとなるのは、式を表現するExprです。
pl.col()で返ってくる型はExprとなっており、Exprを組み合わせた計算式、例えば(pl.col("Embarked") == "S") & (pl.col("Age").is_not_null())も得られる型はExprです。演算子もオーバライドされています。

Exprの一連の流れ (コンテキスト) が最終的には評価・処理されます。
多くは列選択 (select, with_columns) -> 行選択 (filter) -> 集約 (groupby, agg) の順番となります。

Polarsでは、上述の式とコンテキストでデータ変換用のDSLを独自に定義しています。

カラム (col)

col()は、カラムを表すExprが返されます。後述するselectなどと組み合わせて使われます。
引数はカラム名で、正規表現も受け付けます。"*"としたら全てのカラムを表します。
上述の通り、col()で返されたExpr同士やExprとリテラル値の計算も、Exprを返します。

col = pl.col("a")

print(col)
# col("a")

print(type(col))
# <class 'polars.expr.expr.Expr'>

print(type(pl.col("a") + pl.col("b") + 1))
# <class 'polars.expr.expr.Expr'>

カラムの選択 (select, alias, with_columns)

select()は、DataFrameからカラムを選択し、新たなDataFrameを返すメソッドです。
select()の引数には、選択する列をExprのIterableで渡します。ここでもExprなので、pl.col("a") + 1で値の加工も表現できるようになっています。

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

df.select([pl.col("a")])
# shape: (3, 1)
# ┌─────┐
# │ a   │
# │ --- │
# │ i64 │
# ╞═════╡
# │ 1   │
# │ 2   │
# │ 3   │
# └─────┘

df.select([pl.col("a") + 1, pl.col("b") * 2])
# shape: (3, 2)
# ┌─────┬─────┐
# │ a   ┆ b   │
# │ --- ┆ --- │
# │ i64 ┆ i64 │
# ╞═════╪═════╡
# │ 2   ┆ 8   │
# │ 3   ┆ 10  │
# │ 4   ┆ 12  │
# └─────┴─────┘

Exprを実行した場合、結果は元のカラムと同じ名前になります。DataFrame内で同じカラム名は許容されないので、同じカラムに対して複数の式を実行してしまうとエラーとなります。
このエラーはalias()でカラム名を変えることで回避できます。

# NOTE: 2つのカラムはいずれも、カラム名が"a"となる
df.select([
    (pl.col("a") + pl.col("b")),
    (pl.col("a") - pl.col("b")),
])
# DuplicateError: column with name 'a' has more than one occurrences

df.select([
    (pl.col("a") + pl.col("b")).alias("a + b"),
    (pl.col("a") - pl.col("b")).alias("a - b"),
])
# ┌───────┬───────┐
# │ a + b ┆ a - b │
# │ ---   ┆ ---   │
# │ i64   ┆ i64   │
# ╞═══════╪═══════╡
# │ 5     ┆ -3    │
# │ 7     ┆ -3    │
# │ 9     ┆ -3    │
# └───────┴───────┘

もともとDataFrameで持っていた全てのカラムを残して、新たにカラムを追加したい場合は、with_columns()を使います。

df.with_columns([
    (pl.col("a") + pl.col("b")).alias("a + b"),
])
# shape: (3, 3)
# ┌─────┬─────┬───────┐
# │ a   ┆ b   ┆ a + b │
# │ --- ┆ --- ┆ ---   │
# │ i64 ┆ i64 ┆ i64   │
# ╞═════╪═════╪═══════╡
# │ 1   ┆ 4   ┆ 5     │
# │ 2   ┆ 5   ┆ 7     │
# │ 3   ┆ 6   ┆ 9     │
# └─────┴─────┴───────┘

ユニーク数 (n_unique, approx_unique)

ここからselect()で頻出する関数やメソッドを見ていきます。

値のユニークな数をカウントする関数には、n_unique()approx_unique()があります。

n_unique()は全ての値から集計した厳密な結果ですが、approx_unique()は一部の値から近似して算出されます。

titanic_df.select([
    pl.col("Cabin").n_unique().alias("nunique_cabin"),
    pl.approx_unique("Cabin").alias("approx_unique_cabin"),
])
# ┌───────────────┬─────────────────────┐
# │ nunique_cabin ┆ approx_unique_cabin │
# │ ---           ┆ ---                 │
# │ u32           ┆ u32                 │
# ╞═══════════════╪═════════════════════╡
# │ 148           ┆ 148                 │
# └───────────────┴─────────────────────┘

分岐 (when, then, otherwise)

条件式に応じた分岐のための関数およびメソッドとして、when()、then()、otherwise()があります。

when()がtrueのときはthen()が、そうでなければotherwise()が、各行に対して実行されます。then()の後にさらにwhen()を続けることもできるので、3つ以上の分岐も表現できます。

# NOTE: Ageに応じた文字列 ("unknown", "adult", "child") をage_groupにセットする
titanic_df.select(
    [
        pl.col("Age"),
        pl.when(pl.col("Age").is_null())
        .then(pl.lit("unknown"))
        .when(pl.col("Age") >= 18)
        .then(pl.lit("adult"))
        .otherwise(pl.lit("child"))
        .alias("age_group"),
    ]
)
.head(9)
# ┌──────┬───────────┐
# │ Age  ┆ age_group │
# │ ---  ┆ ---       │
# │ f64  ┆ str       │
# ╞══════╪═══════════╡
# │ 22.0 ┆ adult     │
# │ 38.0 ┆ adult     │
# │ 26.0 ┆ adult     │
# │ 35.0 ┆ adult     │
# │ 35.0 ┆ adult     │
# │ null ┆ unknown   │
# │ 54.0 ┆ adult     │
# │ 2.0  ┆ child     │
# │ 27.0 ┆ adult     │
# └──────┴───────────┘

欠損値の扱い

when()でnullが出てきたので、欠損値の扱いを確認しておきます。

Arrow同様、欠損値はデータ型に関わらずnullで表現されます。浮動小数点型については、さらにNaNという値も存在しますが、これは欠損値としては扱われません。
mean()などでもnullは無視されます。

df = pl.DataFrame({
    "a": [1, 2, 3, None],
    "b": [1.0, None, 3.0, np.nan]
})
# ┌──────┬──────┐
# │ a    ┆ b    │
# │ ---  ┆ ---  │
# │ i64  ┆ f64  │
# ╞══════╪══════╡
# │ 1    ┆ 1.0  │
# │ 2    ┆ null │
# │ 3    ┆ 3.0  │
# │ null ┆ NaN  │
# └──────┴──────┘

df.select([
    pl.col("a").count().alias("a_count"),
    pl.col("b").count().alias("b_count"),
    pl.col("a").mean().alias("a_mean"),
    pl.col("b").mean().alias("b_mean"),
])
# ┌─────────┬─────────┬────────┬────────┐
# │ a_count ┆ b_count ┆ a_mean ┆ b_mean │
# │ ---     ┆ ---     ┆ ---    ┆ ---    │
# │ u32     ┆ u32     ┆ f64    ┆ f64    │
# ╞═════════╪═════════╪════════╪════════╡
# │ 4       ┆ 4       ┆ 2.0    ┆ NaN    │
# └─────────┴─────────┴────────┴────────┘

is_null()is_not_null()などを使って行選択できます。

df.filter(pl.col("a").is_not_null())
# ┌─────┬──────┐
# │ a   ┆ b    │
# │ --- ┆ ---  │
# │ i64 ┆ f64  │
# ╞═════╪══════╡
# │ 1   ┆ 1.0  │
# │ 2   ┆ null │
# │ 3   ┆ 3.0  │
# └─────┴──────┘

fill_null()による値埋めや、interpolate()による補間などができます。
fill_null()は引数strategyで固定値以外も指定できます (mean, forwardなど) ので、必要に応じて確認しましょう。

df.select([
    pl.col("a").fill_null(0).alias("a_fill_0"),
    pl.col("a").fill_null(strategy="forward").alias("a_fill_forward"),
    pl.col("a").fill_null(pl.mean("a")).alias("a_fill_mean"),
    pl.col("b").interpolate().alias("b_interpolate"),
    pl.col("b").fill_nan(0).alias("b_fill_nan"),
])
# ┌──────────┬────────────────┬─────────────┬───────────────┬────────────┐
# │ a_fill_0 ┆ a_fill_forward ┆ a_fill_mean ┆ b_interpolate ┆ b_fill_nan │
# │ ---      ┆ ---            ┆ ---         ┆ ---           ┆ ---        │
# │ i64      ┆ i64            ┆ f64         ┆ f64           ┆ f64        │
# ╞══════════╪════════════════╪═════════════╪═══════════════╪════════════╡
# │ 1        ┆ 1              ┆ 1.0         ┆ 1.0           ┆ 1.0        │
# │ 2        ┆ 2              ┆ 2.0         ┆ 2.0           ┆ null       │
# │ 3        ┆ 3              ┆ 3.0         ┆ 3.0           ┆ 3.0        │
# │ 0        ┆ 3              ┆ 2.0         ┆ NaN           ┆ 0.0        │
# └──────────┴────────────────┴─────────────┴───────────────┴────────────┘

型変換 (cast)

型変換はcast()メソッドで行います。nullの値はcastしてもそのままnullになります。

pl.DataFrame({
    "a": [100, 200, None, 300],
}).select([
    pl.col("a").cast(pl.Int64).alias("a_i64"),
    pl.col("a").cast(pl.Float64).alias("a_f64"),
])
# ┌───────┬───────┐
# │ a_i64 ┆ a_f64 │
# │ ---   ┆ ---   │
# │ i64   ┆ f64   │
# ╞═══════╪═══════╡
# │ 100   ┆ 100.0 │
# │ 200   ┆ 200.0 │
# │ null  ┆ null  │
# │ 300   ┆ 300.0 │
# └───────┴───────┘

cast()は、引数strictを持っていて、True (デフォルト) の場合は変換できないとエラーになります。例えば以下のように、ダウンキャストによるオーバーフローはComputeErrorとなります。

titanic_df.select([
    pl.col("PassengerId").cast(pl.Int8).alias("PassengerId_i8"),
]).collect()
# ComputeError: strict conversion from `i64` to `i8` failed for value(s) [128, 129, … 891]; if you were trying to cast Utf8 to temporal dtypes, consider using `strptime`

strict=Falseで変換できなかった場合、値はnullになります。

# NOTE: PassengerIdは1始まりの連番なので、Int8にキャストすると、128番目からnullになる
titanic_df.select(
    [
        pl.col("PassengerId").cast(pl.Int8, strict=False).alias("PassengerId_i8"),
    ]
).head(130).tail(5)
# ┌────────────────┐
# │ PassengerId_i8 │
# │ ---            │
# │ i8             │
# ╞════════════════╡
# │ 126            │
# │ 127            │
# │ null           │
# │ null           │
# │ null           │
# └────────────────┘

ちなみにクエリは遅延評価されるので、以下のように .head(3) とするとPassengerIdの値は1, 2, 3となり、Int8でオーバーフローしないのでエラーにならずに変換されます。

titanic_df.select(
    [
        pl.col("PassengerId").cast(pl.Int8).alias("PassengerId_i8"),
    ]
).head(3)

文字列からCategoricalや数値への変換もcast()でOKです。
数値に変換できない値を含む場合 (以下の例では"not a number") 、strict=Falseでnullを含む変換をしてから、fill_null()で埋めるなどの工夫が必要です。

pl.DataFrame({
    "a": ["10", "2.65", "-20", "-300000.24"],
    "b": ["10", "12", "not a number", "14"],
}).select([
    # NOTE: 以下の変換はComputeError
    # pl.col("a").cast(pl.Int64).alias("a_i64"),
    # pl.col("b").cast(pl.Int8).alias("b_i8"),

    pl.col("a").cast(pl.Float32).alias("a_f32"),
    # strict=Falseでキャストしてから、fill_nullで欠損値を埋める
    pl.col("b").cast(pl.Int8, strict=False).fill_null(0).alias("b_i8"),
])
# shape: (4, 2)
# ┌────────────┬──────┐
# │ a_f32      ┆ b_i8 │
# │ ---        ┆ ---  │
# │ f32        ┆ i8   │
# ╞════════════╪══════╡
# │ 10.0       ┆ 10   │
# │ 2.65       ┆ 12   │
# │ -20.0      ┆ 0    │
# │ -300000.25 ┆ 14   │
# └────────────┴──────┘

数値からBooleanにcast()で変換する場合、0はfalse、それ以外はtrueになります。

pl.DataFrame({
    "a": [-1, 0, 1, 2, None],
    "b": [-1.0, 0, 1.0, 2.0, None],
    "c": ["c1", "c2", "c3", "c3", "c4"],
}).select([
    pl.col("a").cast(pl.Boolean).alias("a_bool"),
    pl.col("b").cast(pl.Boolean).alias("b_bool"),
    # NOTE: 以下の変換はエラー (ArrowErrorException: NotYetImplemented("Casting from LargeUtf8 to Boolean not supported"))
    # pl.col("c").cast(pl.Boolean).alias("c_bool"),
])
# ┌────────┬────────┐
# │ a_bool ┆ b_bool │
# │ ---    ┆ ---    │
# │ bool   ┆ bool   │
# ╞════════╪════════╡
# │ true   ┆ true   │
# │ false  ┆ false  │
# │ true   ┆ true   │
# │ true   ┆ true   │
# │ null   ┆ null   │
# └────────┴────────┘

文字列

文字列カラムへの操作メソッドは .str でアクセスできます。
おおむねPandasと同様なのですが、文字列長の扱いに注意です。lengths()はバイト長を返すので、非ASCII文字以外を含む文字列長はn_chars()で取得します。

pl.DataFrame({
    "a": ["a", "あ", "ab", "abc", None],
}).select([
    pl.col("a").str.lengths().alias("length"),
    pl.col("a").str.n_chars().alias("n_chars"),
    pl.col("a").str.contains("a").alias("contains"),
    pl.col("a").str.starts_with("ab").alias("starts_with"),
    pl.col("a").str.extract(r"a(\w+)", group_index=1).alias("extract"),
    pl.col("a").str.replace_all("ab", "x").alias("replace_all"),
])
# ┌────────┬─────────┬──────────┬─────────────┬─────────┬─────────────┐
# │ length ┆ n_chars ┆ contains ┆ starts_with ┆ extract ┆ replace_all │
# │ ---    ┆ ---     ┆ ---      ┆ ---         ┆ ---     ┆ ---         │
# │ u32    ┆ u32     ┆ bool     ┆ bool        ┆ str     ┆ str         │
# ╞════════╪═════════╪══════════╪═════════════╪═════════╪═════════════╡
# │ 1      ┆ 1       ┆ true     ┆ false       ┆ null    ┆ a           │
# │ 3      ┆ 1       ┆ false    ┆ false       ┆ null    ┆ あ          │
# │ 2      ┆ 2       ┆ true     ┆ true        ┆ b       ┆ x           │
# │ 3      ┆ 3       ┆ true     ┆ true        ┆ bc      ┆ xc          │
# │ null   ┆ null    ┆ null     ┆ null        ┆ null    ┆ null        │
# └────────┴─────────┴──────────┴─────────────┴─────────┴─────────────┘

struct型

以下のscoresのようにカラムの一要素をdictで作ることもできます。この場合struct型となります。
struct型の値の操作は、.structでアクセスできます。

struct型はスキーマを持っており、スキーマと合致した値が無い場合 (以下では{"hoge": "fuga", "bar": 100})、はnullになります。
schemaプロパティで確認すると、Utf8型のsubjectとInt64型のscoreの2つの値を持ちます。

field()で、指定のキーで値にアクセスできます。
また、unnest()で、値をDataFrameのカラムに展開できます。カラム名はキーが使われます。

scores = pl.Series(
    "scores",
    [
        {"subject": "math", "score": 70},
        {"subject": "history", "score": 90},
        {"subject": "physics", "score": 85},
        {"subject": "English", "score": 70},
        {"hoge": "fuga", "bar": 100}
    ],
)
# shape: (5,)
# Series: 'scores' [struct[2]]
# [
# 	{"math",70}
# 	{"history",90}
# 	{"physics",85}
# 	{"English",70}
# 	{null,null}
# ]

scores.struct.schema
# {'subject': Utf8, 'score': Int64}

scores.struct.field("score")
# shape: (5,)
# Series: 'score' [i64]
# [
# 	70
# 	90
# 	85
# 	70
# 	null
# ]

# NOTE: 以下はエラー (StructFieldNotFoundError: bar)
# scores.struct.field("bar")

scores.struct.unnest()
# shape: (5, 2)
# ┌─────────┬───────┐
# │ subject ┆ score │
# │ ---     ┆ ---   │
# │ str     ┆ i64   │
# ╞═════════╪═══════╡
# │ math    ┆ 70    │
# │ history ┆ 90    │
# │ physics ┆ 85    │
# │ English ┆ 70    │
# │ null    ┆ null  │
# └─────────┴───────┘

ウィンドウ関数

mean()やrank()などのExprは、over()でグループ化することもできます。
これによって、SQLでのウィンドウ関数と同等の操作を実装できます。

以下はSQLで書けば、SELECT PassengerId, Pclass, Fare, AVG(Fare) OVER (PARTITION BY Pclass) AS Fare_mean_by_Pclass, RANK() OVER (PARTITION BY Pclass ORDER BY Fare DESC) AS Fare_rank_by_Pclass FROM ...となります。

titanic_df.select(
    [
        pl.col("PassengerId"),
        pl.col("Pclass"),
        pl.col("Fare"),
        pl.col("Fare").mean().over("Pclass").alias("Fare_mean_by_Pclass"),
        pl.col("Fare")
        .rank(descending=True)
        .over("Pclass")
        .alias("Fare_rank_by_Pclass"),
    ]
).head()
# ┌─────────────┬────────┬─────────┬─────────────────────┬─────────────────────┐
# │ PassengerId ┆ Pclass ┆ Fare    ┆ Fare_mean_by_Pclass ┆ Fare_rank_by_Pclass │
# │ ---         ┆ ---    ┆ ---     ┆ ---                 ┆ ---                 │
# │ i64         ┆ i64    ┆ f64     ┆ f64                 ┆ f32                 │
# ╞═════════════╪════════╪═════════╪═════════════════════╪═════════════════════╡
# │ 1           ┆ 3      ┆ 7.25    ┆ 13.67555            ┆ 427.0               │
# │ 2           ┆ 1      ┆ 71.2833 ┆ 84.154687           ┆ 98.0                │
# │ 3           ┆ 3      ┆ 7.925   ┆ 13.67555            ┆ 271.5               │
# │ 4           ┆ 1      ┆ 53.1    ┆ 84.154687           ┆ 123.0               │
# │ 5           ┆ 3      ┆ 8.05    ┆ 13.67555            ┆ 240.0               │
# └─────────────┴────────┴─────────┴─────────────────────┴─────────────────────┘

集約処理 (groupby, agg)

続いて集約処理です。

groupby()でグループ化するカラムを選択 (Iterableで複数指定可) して、agg()で集計する処理を記述します。
agg()でもselect()と同様にExprで表現できますので、Pandasのagg()よりも一貫性があってシンプルな記述ができます。
ただただカラムを選択することもできて (以下の例ではname_list) 、その場合はlist型になります。

titanic_df.groupby("Embarked").agg(
    [
        pl.count().alias("count"),
        (pl.col("Age") < 18).sum().alias("count_under_18"),
        pl.col("Name").alias("name_list"),
        pl.mean("Fare").alias("fare_mean"),
        pl.col("Fare").filter(pl.col("Age") < 18).mean().alias("fare_mean_under_18"),
    ]
).sort("count", descending=True)
# ┌──────────┬───────┬────────────────┬─────────────────────────────┬───────────┬────────────────────┐
# │ Embarked ┆ count ┆ count_under_18 ┆ name_list                   ┆ fare_mean ┆ fare_mean_under_18 │
# │ ---      ┆ ---   ┆ ---            ┆ ---                         ┆ ---       ┆ ---                │
# │ str      ┆ u32   ┆ u32            ┆ list[str]                   ┆ f64       ┆ f64                │
# ╞══════════╪═══════╪════════════════╪═════════════════════════════╪═══════════╪════════════════════╡
# │ S        ┆ 644   ┆ 82             ┆ ["Braund, Mr. Owen Harris", ┆ 27.079812 ┆ 33.722866          │
# │          ┆       ┆                ┆ "Hei…                       ┆           ┆                    │
# │ C        ┆ 168   ┆ 24             ┆ ["Cumings, Mrs. John        ┆ 59.954144 ┆ 25.944279          │
# │          ┆       ┆                ┆ Bradley (Fl…                ┆           ┆                    │
# │ Q        ┆ 77    ┆ 7              ┆ ["Moran, Mr. James", "Rice, ┆ 13.27603  ┆ 20.001786          │
# │          ┆       ┆                ┆ Mast…                       ┆           ┆                    │
# │ null     ┆ 2     ┆ 0              ┆ ["Icard, Miss. Amelie",     ┆ 80.0      ┆ null               │
# │          ┆       ┆                ┆ "Stone, …                   ┆           ┆                    │
# └──────────┴───────┴────────────────┴─────────────────────────────┴───────────┴────────────────────┘

ユーザ定義関数 (map, apply)

DataFrameの各要素に対して、ユーザが定義した関数やラムダ式を適用するメソッドとして、map()apply()が提供されています。
mapはSeries、applyはSeriesの要素を引数に受け取ります。

pl.DataFrame({
    "group": ["a", "a", "b", "b", "b"],
    "value": [1, 2, 3, 4, 5],
}).groupby("group").agg([
    pl.col("value").sum().alias("sum"),
    pl.col("value").map(lambda s: s + 1).alias("value + 1 list"),
])
# ┌───────┬─────┬────────────────┐
# │ group ┆ sum ┆ value + 1 list │
# │ ---   ┆ --- ┆ ---            │
# │ str   ┆ i64 ┆ list[i64]      │
# ╞═══════╪═════╪════════════════╡
# │ b     ┆ 12  ┆ [4, 5, 6]      │
# │ a     ┆ 3   ┆ [2, 3]         │
# └───────┴─────┴────────────────┘

titanic_df.select([
    pl.col("Name").apply(lambda s: "Mrs." in s).alias("is_mrs"),
]).head(2)
# ┌────────┐
# │ is_mrs │
# │ ---    │
# │ bool   │
# ╞════════╡
# │ false  │
# │ true   │
# └────────┘

ラムダ式をmap()やapply()で実行するのは、Polarsがサポートするその他の式構文を使う場合と比較して、遅いという点は注意が必要です。
以下のような簡単なラムダ式でも .str.contains() を使った方が速いことがわかります。

%%timeit
titanic_df.select([
    pl.col("Name").apply(lambda s: "Mrs." in s).alias("is_mrs"),
])
# 482 µs ± 13.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%%timeit
titanic_df.select([
    pl.col("Name").str.contains("Mrs.").alias("is_mrs"),
])
# 384 µs ± 5.83 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

join

DataFrame同士の横方向 (列) の結合は、join()で行います。
引数howには innerleftoutersemianticross を指定できます。

members = pl.DataFrame({
    "id": [1, 2, 3, 4],
    "name": ["Tanaka", "Kudou", "Hirota", "Umeda"],
    "department_id": [1, 2, 1, 3],
})

departments = pl.DataFrame({
    "id": [1, 2],
    "name": ["Sales", "Marketing"]
})

members.join(departments, how="inner", left_on="department_id", right_on="id")
# ┌─────┬────────┬───────────────┬────────────┐
# │ id  ┆ name   ┆ department_id ┆ name_right │
# │ --- ┆ ---    ┆ ---           ┆ ---        │
# │ i64 ┆ str    ┆ i64           ┆ str        │
# ╞═════╪════════╪═══════════════╪════════════╡
# │ 1   ┆ Tanaka ┆ 1             ┆ Sales      │
# │ 2   ┆ Kudou  ┆ 2             ┆ Marketing  │
# │ 3   ┆ Hirota ┆ 1             ┆ Sales      │
# └─────┴────────┴───────────────┴────────────┘

members.join(departments, how="left", left_on="department_id", right_on="id")
# ┌─────┬────────┬───────────────┬────────────┐
# │ id  ┆ name   ┆ department_id ┆ name_right │
# │ --- ┆ ---    ┆ ---           ┆ ---        │
# │ i64 ┆ str    ┆ i64           ┆ str        │
# ╞═════╪════════╪═══════════════╪════════════╡
# │ 1   ┆ Tanaka ┆ 1             ┆ Sales      │
# │ 2   ┆ Kudou  ┆ 2             ┆ Marketing  │
# │ 3   ┆ Hirota ┆ 1             ┆ Sales      │
# │ 4   ┆ Umeda  ┆ 3             ┆ null       │
# └─────┴────────┴───────────────┴────────────┘

concat

DataFrame同士の縦方向 (行) の連結では、concat()を使います。
こちらも引数howにて、verticalhorizontaldiagonalalignを指定可能です。

pl.concat(
    [
        pl.DataFrame(
            {
                "id": [1, 3],
                "name": ["Tanaka", "Hirota"],
            }
        ),
        pl.DataFrame(
            {
                "id": [2, 4],
                "name": ["Kudou", "Nakagawa"],
            }
        ),
    ]
)
# shape: (4, 2)
# ┌─────┬──────────┐
# │ id  ┆ name     │
# │ --- ┆ ---      │
# │ i64 ┆ str      │
# ╞═════╪══════════╡
# │ 1   ┆ Tanaka   │
# │ 3   ┆ Hirota   │
# │ 2   ┆ Kudou    │
# │ 4   ┆ Nakagawa │
# └─────┴──────────┘

pl.concat(
    [
        pl.DataFrame(
            {
                "id": [1, 2, 3, 4],
                "name": ["Tanaka", "Kudou", "Hirota", "Umeda"],
            }
        ),
        pl.DataFrame(
            {
                "department_id": [1, 2, 1, 3],
            }
        ),
    ],
    how="horizontal",
)
# shape: (4, 3)
# ┌─────┬────────┬───────────────┐
# │ id  ┆ name   ┆ department_id │
# │ --- ┆ ---    ┆ ---           │
# │ i64 ┆ str    ┆ i64           │
# ╞═════╪════════╪═══════════════╡
# │ 1   ┆ Tanaka ┆ 1             │
# │ 2   ┆ Kudou  ┆ 2             │
# │ 3   ┆ Hirota ┆ 1             │
# │ 4   ┆ Umeda  ┆ 3             │
# └─────┴────────┴───────────────┘

クエリの実行計画

LazyDataFrameの場合、explain()でクエリの実行計画を確認できます。
https://pola-rs.github.io/polars-book/user-guide/lazy/query_plan/

最適化の有無は引数optimizedで指定します。
optimized=Falseの場合、最初に全列 (PROJECT */12 COLUMNS) が読み込まれます。
一方、optimized=Trueの場合はロードするのは3列になり (PROJECT 3/12 COLUMNS) 、さらにfilterによって行も絞られる (SELECTION: [(col("Sex")) == (Utf8(male))]) ことが読み取れます。

lazy_df = pl.scan_csv("data/titanic.csv")

# 最適化無しの場合
lazy_df.filter(pl.col("Sex") == "male").select(["Survived", "Age"]).groupby("Survived").agg(pl.mean("Age")).explain(optimized=False)
# AGGREGATE
# 	[col("Age").mean()] BY [col("Survived")] FROM
# 	 SELECT [col("Survived"), col("Age")] FROM
#   FILTER [(col("Sex")) == (Utf8(male))] FROM
#     CSV SCAN data/titanic.csv
#     PROJECT */12 COLUMNS

# 最適化有りの場合
lazy_df
    .filter(pl.col("Sex") == "male")
    .select(["Survived", "Age"])
    .groupby("Survived")
    .agg(pl.mean("Age"))
    .explain(optimized=True)
# AGGREGATE
# 	[col("Age").mean()] BY [col("Survived")] FROM
# 	FAST_PROJECT: [Survived, Age]
#     CSV SCAN data/titanic.csv
#     PROJECT 3/12 COLUMNS
#     SELECTION: [(col("Sex")) == (Utf8(male))]

他との互換性

pl.DataFrameからの変換

to_X()でpl.DataFrameから他の形式へ変換するメソッドが用意されています。

to_pandas()でpandas.DataFrameに変換されます。
引数use_pyarrow_extension_array=True とすると、numpy配列の代わりにPyArrow配列で保持されるようになり、その場合はpl.DataFrameからの値コピーは発生しません。

pandas_numpy_df = titanic_df.to_pandas()
pandas_numpy_df.info()
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 891 entries, 0 to 890
# Data columns (total 12 columns):
#  #   Column       Non-Null Count  Dtype  
# ---  ------       --------------  -----  
#  0   PassengerId  891 non-null    int64  
#  1   Survived     891 non-null    int64  
#  2   Pclass       891 non-null    int64  
#  3   Name         891 non-null    object 
#  4   Sex          891 non-null    object 
#  5   Age          714 non-null    float64
#  6   SibSp        891 non-null    int64  
#  7   Parch        891 non-null    int64  
#  8   Ticket       891 non-null    object 
#  9   Fare         891 non-null    float64
#  10  Cabin        204 non-null    object 
#  11  Embarked     889 non-null    object 
# dtypes: float64(2), int64(5), object(5)
# memory usage: 83.7+ KB

pandas_arrow_df = titanic_df.to_pandas(use_pyarrow_extension_array=True)
print(pandas_arrow_df.info())
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 891 entries, 0 to 890
# Data columns (total 12 columns):
#  #   Column       Non-Null Count  Dtype                
# ---  ------       --------------  -----                
#  0   PassengerId  891 non-null    int64[pyarrow]       
#  1   Survived     891 non-null    int64[pyarrow]       
#  2   Pclass       891 non-null    int64[pyarrow]       
#  3   Name         891 non-null    large_string[pyarrow]
#  4   Sex          891 non-null    large_string[pyarrow]
#  5   Age          714 non-null    double[pyarrow]      
#  6   SibSp        891 non-null    int64[pyarrow]       
#  7   Parch        891 non-null    int64[pyarrow]       
#  8   Ticket       891 non-null    large_string[pyarrow]
#  9   Fare         891 non-null    double[pyarrow]      
#  10  Cabin        204 non-null    large_string[pyarrow]
#  11  Embarked     889 non-null    large_string[pyarrow]
# dtypes: double[pyarrow](2), int64[pyarrow](5), large_string[pyarrow](5)
# memory usage: 119.3 KB

NumPy

PolarsのExprでは、NumPyのuniversal functions (ufunc)をサポートしてます。

pl.DataFrame({"a": [1, 2, 3, None]}).select([
    np.sqrt(pl.col("a")).alias("sqrt"),
    np.exp(pl.col("a")).alias("exp"),
])
# ┌──────────┬───────────┐
# │ sqrt     ┆ exp       │
# │ ---      ┆ ---       │
# │ f64      ┆ f64       │
# ╞══════════╪═══════════╡
# │ 1.0      ┆ 2.718282  │
# │ 1.414214 ┆ 7.389056  │
# │ 1.732051 ┆ 20.085537 │
# │ null     ┆ null      │
# └──────────┴───────────┘

まとめ

Polarsの特徴と使い勝手についてざっと見てきましたが、最後に所感です。
今回、パフォーマンス面などは見てないですが、Pandasと比較してAPIが素敵なので、今から分析や開発のタスクを始めるならPandasよりもPolarsを使いたい、という気持ちです。

  • APIがシンプルで構造的
    • タイプ量は増えがちだが、Exprやselect、filter、groupbyなどで統一的に記述できる
      • SQLを書ける人にはPandasよりも馴染みやすいと思います
    • インデックスを持たないなどPandasのネガティブな仕様も解消されている一方で、ネーミングなどはPandasに近いのでPandas使いもとっつきやすい
  • 型、スキーマに厳密
    • データに起因した予想外の振る舞い (Pandasの整数列に空値が含まれていた場合に型がfloatになる...など) が起きづらい
    • 欠損値は全てnull扱い、基本的にObjectを使う必要がない、などもポイント
    • Arrowを基本的に踏襲している恩恵
  • 最適化処理もうまく隠蔽されている
    • 書き方や順序がブレても、最終的な評価では同じ実行計画になることが期待できる
    • 式やコンテキストの遅延評価の恩恵

参考

Apache Arrow and the "10 Things I Hate About pandas" - Wes McKinney
超高速…だけじゃない!Pandasに代えてPolarsを使いたい理由 - Qiita
Polars, 旬の13のお役立ち機能 - Qiita
Rust製高速データフレームライブラリ、Polarsを試す | gihyo.jp

Discussion