PandasからPolarsに移行して混乱するケース3選

2023年10月31日 16:13

Pandasに慣れた人がPolarsでハマりやすそうなケースをまとめました。
PolarsはPandasと似たメソッドが多いので、Pandasに慣れているとドキュメントを参照してすぐに簡単な操作ができるようになっています。
ただ、ドキュメントには複雑な操作例があまり載っていないこともあり、少し込み入ったことしようとするとPandasとの差異が重なって混乱しやすい気がします。
ChatGPTに聞いてもPandasとの混乱が見られる場合が経験上結構あります。

1. ある列の値でグループ化して、別の列を加工する

df = df.sort("group")
df = df.with_columns(
    pl.col("a").shift(1).over("group").alias("shifted a"),
    pl.col("b").forward_fill().over("group").alias("filled b"),
)

Pandasだとgroupbyを使うと思います。Polarsにもgroup_byがあるのでついgroup_byから考えてしまいがちです。Polarsでは、overを使うことである列の値ごとにgroup化して別の列を処理できます。
注意点として、事前にソートしておく必要があります。
上の例は、group列でグループ化してa列やb列を加工しています。aliasで元の列名を指定すれば、元の列を置換します。
sumなどのagg処理を行う場合はもちろんgroup_byです。

2. ある列の値を条件として、別の列に適用する関数を変える

df = df.with_colums(
    pl.when(pl.col("year") == "next year")\
            .then(pl.("date") + timedelta(days=365))\
                .otherwise(pl.col("date"))\
                    .alias("date"),
)

Polarsドキュメントのwhenの例には関数の適用はありませんが、whenで書けます。thenやotherwiseの中で処理を行えます。applyとかmap_elementとか使う必要はありません。
上の例は、year列の値が"next year"のときだけ、date列の値を一年未来にするというものです。（こういうふざけたワイルドデータにたまに遭遇します）

3. 無限大をNullに置換する

df = pl.DataFrame({
    "floats": [1.0, 2.0, np.inf, -np.inf, 5.0],
    "ints": [1, 2, 3, 4, 5],
    "strings": ["a", "b", "c", "d", "e"]
})

numeric_cols = pl.selectors.by_dtype(pl.NUMERIC_DTYPES)
df = df.with_columns(
    pl.when(numeric_cols.is_infinite()).then(None).otherwise(numeric_cols).keep_name()
)

Pandasだと df.replace({np.inf: np.nan}) で一発ですが、Polarsのreplaceは全然違う（列を丸ごと別のSeriesに置き換える）ので、when-then-otherwiseを使います。
また、is_infiniteというメソッドが用意されているのは有難いですが、適用可能なデータ型が限られているので、pl.selectors.by_dtypeを使って見る列を絞ります。（全ての列が数値型であることが分かっている場合は、numeric_cols を取得せずに pl.all() でもOKです。）
さらに、最後にkeep_nameしておくか、対象のカラムの数だけaliasを指定しないといけないという3重トラップです。

[番外編] Pandasだと簡単なのにPolarsでは少しややこしくなるケース

列全体に値を代入する

df = df.with_columns(
    pl.lit("apple").alias("a")
)

Pandas だと df["a"] = "apple" で済みますが、少し複雑になります。
df = df.replace("a", pl.Series(name="a", values=["apple"] * len(df))) などを紹介している記事もありますが、pl.litを使った方がスマートに書けます。

データフレーム全体のsumを取る

df.sum().sum(axis=1)[0]

Pandasでは df.sum().sum() ですが、Polarsではsumのaxisの指定と、要素の明示的な取得が必要です。

データフレーム全体に前方補間をかける

df = df.select(pl.all().forward_fill())

Pandasでは df.ffill() です。
forward_fillのようにpolars.DataFrameには存在せず、polars.Exprのみに実装されているメソッドが多くあります。このようなメソッドの処理をデータフレーム全体に行いたい場合は、selectとpl.all()を使います。

データフレームの各要素が無限大かどうかを取得する

df.select(pl.selectors.by_dtype(pl.NUMERIC_DTYPES).is_infinite())

Pandasだとstr型の列が含まれていたとしても df == np.inf や df.isin([np.inf, -np.inf]) で済みますが、Polarsで同様のことをすると Compute Error になります。
先述の例と同じく、is_infinite と selectors を組み合わせて対応します。

以上です。思いついたら都度追記します。