はじめに
特に pandas ユーザーが Polars に入門する際の最初の障壁はエクスプレッションではないでしょうか(僕はそうでした)。
僕が初めて Polars に入門した際に、pandas での
df_filtered = df[df["col1"] == "A"]
みたいな処理が Polars では
df_filtered = df.filter(pl.col("col1") == "A")
と書かなきゃいけないよ、という情報を見つけたことで出鼻を挫かれ、しばらく入門をやめました。
- pandas っぽい速いライブラリと聞いたけど pandas ぽく書けないのかもしれない
-
pl.col("col1")
って何、学習コスト高そう
と考えたのが理由です。今日もタイムラインから Polars に挫折した人の嘆きが聞こえて来そうです。
そこで、本記事では Polars のエクスプレッションに注目し、特に pandas から Polars の移行を検討している人に対して「エクスプレッションって案外簡単!」という印象を持ってもらい、無理なく Polars に入門してもらうことを目標としています。
一方で、既に Polars を使いこなしているぜ、という人に対しても何かしらかの気づきがあれば幸いです。間違いのご指摘も歓迎します。
また、注意点として、この記事の目的が「エクスプレッションの基本的な理解」のため、API や機能の紹介は最小限に留めています。そのため「具体的に〜な処理はどうしたら良い?」がこの記事で解決される可能性は低いです。それでもエクスプレッションという Polars の重要なコンセプトを理解する一助になれば非常に嬉しく思います。
前置きが長くなりましたが、この記事は Polars Advent Calendar 2023 の最終日、25 日目です!!素晴らしい記事が沢山あるので、他の記事も是非読んでみてください!
Polars エクスプレッションとは
Polars エクスプレッション(polars.Expr
)をざっくり説明すると、
データフレームを処理するために、式(Context)に渡して使う、データ操作の流れを表現するオブジェクト
のような感じだと思います。冒頭の実装例に関しては以下のようになります。
この例では pl.col("col1") == "A"
がエクスプレッションであり式 filter
に対し、操作の流れ( col1
が "A"
であるようなレコードを抽出する)を渡しているイメージです。
Polars エクスプレッションとは
- データフレーム(
polars.DataFrame
)を操作するために使う - 式(Context)と呼ばれるデータフレームのメソッドに引数として渡す
- データ操作の流れを表現するオブジェクト
ここでエクスプレッションはあくまで「操作の流れ」のみを記述する点に注意しましょう。そのため、エクスプレッションが操作後のデータオブジェクトそのものを返すことはありません。実際にデータを処理するのは式(Context)であって、エクスプレッションは式に対してデータ操作の「実行計画書」や「操作手順書」のようなものを渡しているイメージです。
また、この基本操作で式はデータフレームオブジェクト(polars.DataFrame
)を返却します。Inplace 的な操作やその他の様々なデータ操作インターフェースが Polars には提供されていますが、本記事ではこの「式+エクスプレッション」で表現可能な形式についてのみ取り扱います。
式(Context)
Polars で一般的に使う式(Context)は以下です 1。
- シリーズの選択:
df.select(…)
df.with_columns(…)
- データフレームのフィルタリング:
df.filter(…)
- Group by / Aggregation:
df.group_by(…).agg(…)
Polars ではひとまずこれらの式の使い方が分かっていれば多くのデータ操作ができます。Polars に入門するのであればこれらの 4 つの式と「エクスプレッションがそれぞれどのように処理されるか」だけ押さえておけば十分でしょう。
…実は Polars はある程度 pandas-like な構文をサポートしています2。とりあえず Polars を動かす目的であれば pandas みを残しながら最低限の学習で入門できるのでは、という気持ちもわかりますが、後述する並列化の恩恵を受けにくいとされている 3 ため、筆者は基本的に非推奨という立場です。色んな処理を可読性高く書けるというメリットもあるので是非式を使いましょう!
各式(Context)の使い方
それでは、各式がエクスプレッションを受け取って具体的にどのような処理をするのか見ていきましょう。
ここからは以下のようなシンプルなデータフレームで説明します 4。
import numpy as np
import polars as pl
df = pl.DataFrame(
{
"nrs": [1, 2, 3, None, 5],
"names": ["foo", "ham", "spam", "egg", None],
"random": np.random.rand(5),
"groups": ["A", "A", "B", "C", "B"],
}
)
print(df)
shape: (5, 4)
┌──────┬───────┬──────────┬────────┐
│ nrs ┆ names ┆ random ┆ groups │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ f64 ┆ str │
╞══════╪═══════╪══════════╪════════╡
│ 1 ┆ foo ┆ 0.877225 ┆ A │
│ 2 ┆ ham ┆ 0.439351 ┆ A │
│ 3 ┆ spam ┆ 0.59223 ┆ B │
│ null ┆ egg ┆ 0.795323 ┆ C │
│ 5 ┆ null ┆ 0.238454 ┆ B │
└──────┴───────┴──────────┴────────┘
シリーズの選択: df.select(…)
df.with_columns(…)
この式では、データフレームの列に対してエクスプレッションを適用します。前述の通り「ある列に対する一連の処理の流れ」であるエクスプレッションを受け取り、そのとおりに処理を行った結果をデータフレームとして返します。
select
式と with_columns
式は似ていますが、その結果の返し方に違いがあります。
df.select(…)
select
式は引数として渡されたエクスプレッションに従って特定の列を処理し、その列のみを返します。
一番シンプルな例は「指定した列を返すエクスプレッション col
」です。
# names 列を抽出
df_names = df.select(pl.col("names"))
print(df_names)
shape: (5, 1)
┌───────┐
│ names │
│ --- │
│ str │
╞═══════╡
│ foo │
│ ham │
│ spam │
│ egg │
│ null │
└───────┘
col
は指定列そのものを返すエクスプレッションのため、あまりデータを操作している感はありませんが、多くのケースで起点となる重要なエクスプレッションです。
また、結果が 1 列しかありませんが shape が (5, 1)
となっており返り値がデータフレームである点に注意してください 5。
エクスプレッションのインターフェースはかなり抽象化されており、リテラルや別のエクスプレッションなどとの演算評価の多くをカバーしています。「この書き方も行けるかな?」は割といけることが多いです。
# nrs 列に random 列の平均値と 10 を足す
df_result = df.select(pl.col("nrs") + pl.col("random").mean() + 10)
print(df_result)
shape: (5, 1)
┌───────────┐
│ nrs │
│ --- │
│ f64 │
╞═══════════╡
│ 11.588517 │
│ 12.588517 │
│ 13.588517 │
│ null │
│ 15.588517 │
└───────────┘
ここで、pl.col("nrs")
が nrs
列全体を返すので長さ 5 のシリーズであるのに対し、pl.col("random").mean()
や整数リテラル 10
はスカラー値です。これらのスカラー値はブロードキャストされ、長さ 5 ですべてが同じ値であるようなシリーズとして計算されます。
select
式にはカンマ ,
区切りで複数のエクスプレッションを渡せます。一連のエクスプレッションごとに列が計算され、渡したエクスプレッションの数だけ結果の列が増えていきます。
df_result = df.select(
pl.col("names"),
pl.col("nrs") * 2,
pl.col("groups")
)
print(df_result)
shape: (5, 3)
┌───────┬──────┬────────┐
│ names ┆ nrs ┆ groups │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ str │
╞═══════╪══════╪════════╡
│ foo ┆ 2 ┆ A │
│ ham ┆ 4 ┆ A │
│ spam ┆ 6 ┆ B │
│ egg ┆ null ┆ C │
│ null ┆ 10 ┆ B │
└───────┴──────┴────────┘
df.with_columns(…)
一方で with_columns
式は、引数として渡されたエクスプレッションに従って特定の列を処理し、データフレームすべてを返します。先程の実装を、select
から with_columns
に変更してみましょう。
# nrs 列を元の値に random 列の平均値と 10 を足した値に更新
df_result = df.with_columns(pl.col("nrs") + pl.col("random").mean() + 10)
print(df_result)
shape: (5, 4)
┌───────────┬───────┬──────────┬────────┐
│ nrs ┆ names ┆ random ┆ groups │
│ --- ┆ --- ┆ --- ┆ --- │
│ f64 ┆ str ┆ f64 ┆ str │
╞═══════════╪═══════╪══════════╪════════╡
│ 11.588517 ┆ foo ┆ 0.877225 ┆ A │
│ 12.588517 ┆ ham ┆ 0.439351 ┆ A │
│ 13.588517 ┆ spam ┆ 0.59223 ┆ B │
│ null ┆ egg ┆ 0.795323 ┆ C │
│ 15.588517 ┆ null ┆ 0.238454 ┆ B │
└───────────┴───────┴──────────┴────────┘
nrs
列が select
の結果と同じ値に更新された状態でデータフレーム全体が返されたことがわかります。このように with_columns
式では、デフォルトで「基準とする列(今回は nrs
列)を更新する」ような動作になります。
特徴量エンジニアリング等の文脈で、新しい列を追加する場合にはエクスプレッションの alias
メソッドを用いて nrs
と異なる列名をアサインします。
# nrs 列に random 列の平均値と 10 を足した値を feature 列として追加
df_result = df.with_columns(
(pl.col("nrs") + pl.col("random").mean() + 10).alias("feature")
)
print(df_result)
shape: (5, 5)
┌──────┬───────┬──────────┬────────┬───────────┐
│ nrs ┆ names ┆ random ┆ groups ┆ feature │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ f64 ┆ str ┆ f64 │
╞══════╪═══════╪══════════╪════════╪═══════════╡
│ 1 ┆ foo ┆ 0.877225 ┆ A ┆ 11.588517 │
│ 2 ┆ ham ┆ 0.439351 ┆ A ┆ 12.588517 │
│ 3 ┆ spam ┆ 0.59223 ┆ B ┆ 13.588517 │
│ null ┆ egg ┆ 0.795323 ┆ C ┆ null │
│ 5 ┆ null ┆ 0.238454 ┆ B ┆ 15.588517 │
└──────┴───────┴──────────┴────────┴───────────┘
alias
の返却値もエクスプレッションのため、例えばこのように書いても結果は一緒です。
# nrs 列に random 列の平均値と 10 を足した値を feature 列として追加
df_result = df.with_columns(
pl.col("nrs").alias("feature") + pl.col("random").mean() + 10
)
print(df_result)
shape: (5, 5)
┌──────┬───────┬──────────┬────────┬───────────┐
│ nrs ┆ names ┆ random ┆ groups ┆ feature │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ f64 ┆ str ┆ f64 │
╞══════╪═══════╪══════════╪════════╪═══════════╡
│ 1 ┆ foo ┆ 0.877225 ┆ A ┆ 11.588517 │
│ 2 ┆ ham ┆ 0.439351 ┆ A ┆ 12.588517 │
│ 3 ┆ spam ┆ 0.59223 ┆ B ┆ 13.588517 │
│ null ┆ egg ┆ 0.795323 ┆ C ┆ null │
│ 5 ┆ null ┆ 0.238454 ┆ B ┆ 15.588517 │
└──────┴───────┴──────────┴────────┴───────────┘
また、select
式と同様に、複数のエクスプレッションを渡すことで、同時に複数の列を更新したり、新しい列として追加できます。
データフレームのフィルタリング: df.filter(…)
この式では、データフレーム全体に対してエクスプレッションとして渡された条件に合致する行のフィルタリングを行います。すなわち、この式に渡すエクスプレッションは「どのレコードを抽出するか」が表現されるべきであり、Boolean
型のシリーズを返す形式である必要があります。その点で、エクスプレッションの書き方が select
や with_columns
とは異なることに注意してください。
ブールシリーズを返すエクスプレッション
Boolean
型のシリーズを返すエクスプレッションを理解するために、冒頭に登場した実装(チョット違う)で filter
を select
としてみましょう。
# groups 列が "A" かどうかの評価結果を計算
df_result = df.select(pl.col("groups") == "A")
print(df_result)
shape: (5, 1)
┌────────┐
│ groups │
│ --- │
│ bool │
╞════════╡
│ true │
│ true │
│ false │
│ false │
│ false │
└────────┘
select
式の節で「エクスプレッションは多くの演算評価をカバーしている」と説明しましたが、比較演算子の評価結果は上の結果のように Boolean
型になります。単純に「groups
列の値が "A"
と等しいなら true
そうでないなら false
」という評価が行われると考えて問題ありません。
つまり pl.col("groups") == "A"
というエクスプレッションは元のデータフレームの同じ長さを持つブールシリーズを返すエクスプレッションであり、filter
式に渡せる形式になっているのです。
# groups が "A" のようなレコードをデータフレームから抽出
df_filtered = df.filter(pl.col("groups") == "A")
print(df_filtered)
shape: (2, 4)
┌─────┬───────┬──────────┬────────┐
│ nrs ┆ names ┆ random ┆ groups │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ f64 ┆ str │
╞═════╪═══════╪══════════╪════════╡
│ 1 ┆ foo ┆ 0.877225 ┆ A │
│ 2 ┆ ham ┆ 0.439351 ┆ A │
└─────┴───────┴──────────┴────────┘
複数条件を指定したい場合は AND/OR の条件をビット演算子 &
|
で指定します。
# nrs が 1 より大きく groups が "A" であるようなレコードをデータフレームから抽出
df_filtered = df.filter((pl.col("nrs") > 1) & (pl.col("groups") == "A"))
print(df_filtered)
shape: (1, 4)
┌─────┬───────┬──────────┬────────┐
│ nrs ┆ names ┆ random ┆ groups │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ f64 ┆ str │
╞═════╪═══════╪══════════╪════════╡
│ 2 ┆ ham ┆ 0.439351 ┆ A │
└─────┴───────┴──────────┴────────┘
この際、&
|
の優先度が他の比較演算子より高いため、各条件を ()
する必要があることに注意してください。
なお、filter
式に複数のエクスプレッションを渡した場合はそれぞれの条件の AND を取ったものと等しくなります(なので上の実装と結果は同じになります)。
# 複数のエクスプレッションを filter に渡す
df_filtered = df.filter(pl.col("nrs") > 1, pl.col("groups") == "A")
print(df_filtered)
shape: (1, 4)
┌─────┬───────┬──────────┬────────┐
│ nrs ┆ names ┆ random ┆ groups │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ f64 ┆ str │
╞═════╪═══════╪══════════╪════════╡
│ 2 ┆ ham ┆ 0.439351 ┆ A │
└─────┴───────┴──────────┴────────┘
pandas のブールインデックス参照との違い
この Polars における抽出ロジックは一見 pandas におけるブールインデックス参照と似ています。
ブールインデックス参照は pandas において df["groups"] == "A"
の評価結果がブール型のデータシリーズ(ブールベクトル)となり、ブールベクトルをそのまま df[ブールベクトル]
のように参照に用いるというフィルタリング記述で、様々な場面で用いられます 6。
# pandas ブールインデックス参照
df_filtered = df[df["groups"] == "A"]
print(df_filtered)
nrs names random groups
0 1.0 foo 0.877225 A
1 2.0 ham 0.439351 A
Polars のフィルタリング操作との大きな違いは、pandas のブールインデックス参照は df["groups"] == "A"
において一度データ全体の評価を計算した後、ブールベクトルのオブジェクトを生成している点です。Polars では前述のように式に渡されたエクスプレッション自体がデータシリーズのオブジェクトを返さないため、同じブール系列を用いたフィルタリングであっても Polars はメモリや計算効率の面で優れていると言えるでしょう。
Group by / Aggregation: df.group_by(…).agg(…)
最後に Group by / Aggregation の式です。これまでの式との違いは「エクスプレッションを渡すのは基本的に agg
式であり、これは GroupBy
オブジェクトのメソッド」であることです(今までの式たちはデータフレームのオブジェクトでした)。group_by
にもエクスプレッションは渡せなくもないですが、基本的な使い方としては agg
式に「集計したグループごとにどんな処理をするか」を表現するエクスプレッションを渡すことになるでしょう。
ひとまず最もシンプルな col
エクスプレッションを渡してみます。
df_agg = df.group_by("groups").agg(pl.col("nrs"))
print(df_agg)
shape: (3, 2)
┌────────┬───────────┐
│ groups ┆ nrs │
│ --- ┆ --- │
│ str ┆ list[i64] │
╞════════╪═══════════╡
│ B ┆ [3, 5] │
│ C ┆ [null] │
│ A ┆ [1, 2] │
└────────┴───────────┘
同じ groups
の値を持つ nrs
の組がリスト形式で格納されています。 col
に繋げて特定の操作のエクスプレッションを記述することは、これらのリスト内の値にその処理を行うことに等しいです。シンプルな集計処理として合計を計算してみましょう。
df_agg = df.group_by(pl.col("groups")).agg(pl.col("nrs").sum())
print(df_agg)
shape: (3, 2)
┌────────┬─────┐
│ groups ┆ nrs │
│ --- ┆ --- │
│ str ┆ i64 │
╞════════╪═════╡
│ C ┆ 0 │
│ B ┆ 8 │
│ A ┆ 3 │
└────────┴─────┘
無事リストの合計が得られました。
ここで、データの順番が変わっていることに気が付きます。Polars はいくつかの場面で「データの順番を保証しない代わりに計算効率を上げる」ような実装オプションを提供しています。データ順を変えたくない場合は group_by
式(や、その他のメソッド)の引数に maintain_order=True
を渡すことで解決します。
group_by/agg
式にも複数のエクスプレッションを渡せます。以下の実装や出力を眺めてみることで、だんだんエクスプレッションのイメージが湧いてきませんか?
df_agg = df.group_by("groups").agg(
pl.col("nrs"),
pl.col("random").mean(),
pl.col("names").str.concat()
)
print(df_agg)
shape: (3, 4)
┌────────┬───────────┬──────────┬─────────┐
│ groups ┆ nrs ┆ random ┆ names │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ list[i64] ┆ f64 ┆ str │
╞════════╪═══════════╪══════════╪═════════╡
│ A ┆ [1, 2] ┆ 0.658288 ┆ foo-ham │
│ C ┆ [null] ┆ 0.795323 ┆ egg │
│ B ┆ [3, 5] ┆ 0.415342 ┆ spam │
└────────┴───────────┴──────────┴─────────┘
もっとエクスプレッションについて詳しく調べたい方は以下を参照してください。
-
Polars User guide(Context)
- 式(Context)についてのコンセプトが纏まっています
-
Polars User guide(Expressions)
- エクスプレッションについてのコンセプトが纏まっています
-
Polars API reference
- API リファレンスです
コラム
書いたものの本題からは逸れるかなぁと思った節です。
メソッドチェーンについて
エクスプレッションの重要な性質として「エクスプレッションはエクスプレッションのメソッドであり、エクスプレッションを返す」が挙げられます。この性質により、ある処理のエクスプレッション expr1
の後に expr2
expr3
… というふうに処理を複数繋げたいときには、
df.select(expr1.expr2.expr3...)
のように記述できます。なぜなら、expr1
の返却値が polars.Expr
オブジェクトでありメソッドとして expr2
を持つため expr1.expr2
のように直接 expr2
を呼び出せるためです。expr1.expr2
の返却値ももちろん polars.Expr
オブジェクトですから、さらに expr3
を直接呼び出しています。このように一連の処理をドット .
等で繋げて記述する方法はメソッドチェーンと呼ばれます。
Polars ではメソッドチェーンを改行して表記する記述方法がよく用いられます 7。Python の記述ルールではカッコ (){}[]
内であれば自由に改行できるため 8、見やすい形にプログラムを整形できます。
# こんなかんじです
df_selected = df.select(
expr1
.expr2
.expr3
)
また、式たち(select
with_columns
filter
group_by(...).agg
)も同様に「データフレームのメソッドでありデータフレームを返す」ので、やりたい処理のすべてを式を組み合わせて一つのメソッドチェーンで書けたりもします。
# 複数の操作をワンライナーで実現する例
df_result = df.with_columns(pl.col("feature1").add(pl.col("feature2")).alias("f1_plus_f2")).filter(pl.col("f1_plus_f2") > 5).head(5))
…とても横長な記述となってしまいました。
このようなケースでも改行を駆使することで可読性の向上が見込めます。
df_result = (
df
.with_columns(
pl.col("feature1")
.add(pl.col("feature2"))
.alias("f1_plus_f2")
)
.filter(
pl.col("f1_plus_f2") > 5
)
.head(5)
)
単純な可読性だけでなく、処理ごとにコメントを追加したり、一部を変更するような試行錯誤も行いやすいのでオススメです!
ここで with_columns
や filter
といった式を改行した頭から記述するために、コード全体を ()
で囲んでいる点に注意してください。囲まない場合は以下のような記述となり、これもよく見かけます。
df_result = df.with_columns(
pl.col("feature1")
.add(pl.col("feature2"))
.alias("f1_plus_f2")
).filter(
pl.col("f1_plus_f2") > 5
).head(5)
Polars エクスプレッションによって享受できる並列化
Polars はクソ速い(Blazingly Fast な)データフレームライブラリとして知られていますが、理由としてよく挙げられるのは以下の 3 点です。
- バックエンドがコンパイラ言語である Rust で書かれている
- 遅延評価を提供する API、クエリ最適化
- マルチコアでの並列処理を可能な限り利用する
式とエクスプレッションを用いた記述は特に 3 つ目の並列処理と密接に関わっています 9。
例えば各式の紹介で「複数のエクスプレッションを渡した際の挙動」を取り上げましたが、それら複数の一連のエクスプレッションはすべて並列に処理されます 10。少々学習コストを支払ってでも Polars の式+エクスプレッションの表記に慣れておくメリットがここにあります。もちろん特に工夫をしなくても十分高速化が見込めるライブラリではありますが、処理が長く、複雑になるほど並列化の恩恵が大きくなり、処理時間の短縮が見込めます。
Polars の生みの親 Ritchie Vink 氏はこの並列化を「Embarrassingly parallel(恥ずかしいほどの並列化)」と評しています。彼のブログ(Polars オフィシャルサイトに転載)I wrote one of the fastest DataFrame libraries には Polars に実装されている並列化処理のコンセプトやハードウェアレベルでの高速化へのこだわりを読み取ることができるので是非読んでみてください。
おわりに
Polars の API は非常に親切に実装されており、エクスプレッションを使わなくても記述できる処理も多いですが、Polars が輝くのはやはり「大規模なデータの処理パフォーマンスを改善したい」ようなケースだと思います。そのためにはより「Polars 的」な記述で並列化の恩恵を最大化するのが重要だと思います。
また、この記事では言及していませんが LazyAPI を用いた遅延評価・クエリ最適化を活用するのも効果が期待できます。
まだまだ情報の少ない Polars ですが、入門する人が増え、沢山の情報が共有されるようコミュニティを盛り上げていきましょう!
-
https://pola-rs.github.io/polars/user-guide/concepts/contexts/ ↩
-
この辺も将来のアップデートで克服され「pandas っぽい書き方でも爆速!」みたいなことはあるかもしれません ↩
-
Polars のデータフレームは
print
関数でもきれいに表示(pretty-print)されます ↩ -
データシリーズへの変換が必要な場合はデータフレームの
to_series
メソッドを用います。単純にシリーズを取得するのみであればデータフレームのget_column
メソッドも存在しますがこのメソッドにはエクスプレッションが渡せず、最適化の恩恵が受けられない可能性があります ↩ -
最近では
df.query
のほうが人気? ↩ -
公式で明示されているわけではありませんが、ドキュメントにも Discord コミュニティにもよく登場します ↩
-
文字列を横切るような改行はできません ↩
-
LazyAPI を活用する上でも大事なんじゃないかと睨んでいます。参考: Polars と遅延評価 ↩
-
どの程度並列化されるかは使用するシステムに依存します。また、1 つ目のエクスプレッションで作られた列を 2 つ目のエクスプレッションで使用するような場合などは並列化されません ↩