この記事はPolars Advent Calendar 2023の21日目の記事です！
正直こんなの載せるまでもないでしょ、とか色々あるかもしれませんが大目にみていただけると幸いです！

はじめに

Polarsは公式ドキュメントを見ながら実装することが多いのですが、公式ドキュメント見てるとこんな便利な関数あったの？ってパターンが多いのでお気に入りのものを紹介していきます。特にpandasにはこんなのなかったよな？ってものを載せます！この記事はPolars=0.20に基づいて書いていきます。

便利関数一覧

データ読み込み・書き出し

scan_XXX()とsink_XXX()関数です。特にこのsinkが素晴らしいです。　こちらに
レファレンス置いておきます。ここから書いてある内容を引用すると、

This allows streaming results that are larger than RAM to be written to disk.

とのこと。メモリに乗っからなくてもデータ書き出しできるなんてすごいですよね...
データはcsvでもらうけど全部parquetにしたいなって時多いですよね。とんでもサイズcsvでメモリに乗っからないとしても

lf = pl.scan_csv("/path/to/my_larger_than_ram_file.csv")  
lf.sink_parquet("out.parquet")

でいけてしまうんです。便利ですよねえ。

group_byの結果をgroupごとに見る

EDAしてる時にgroup_byして各groupについてどういうデータになってるか見たいな、って時ありませんか？僕はあります。
group_byのレファレンスを見ると

for name, data in df.group_by("a"):  
    print(name)
    print(data)

こうすればできるよ、って書いていますが少し面倒ですよね。そんな時はpartition_by()使いましょう！

df.partition_by("a")

これで上記と同じ結果を得られるんです。これを使わない手はないですよね。

ウィンドウ関数

over()関数です。これをうまく使いこなせるようになると一気にデータ加工が楽になります！
レファレンスによると

Compute expressions over the given groups.

This expression is similar to performing a group by aggregation and joining the result back into the original DataFrame.

The outcome is similar to how window functions work in PostgreSQL.

とあります。groupbyしてデータ加工みたいなことをできるよ、という関数です。この関数のいいところは計算されたグループはキャッシュされ、異なるウィンドウ関数間で共有されます。追加コストがないのです。

説明が難しいのでこれの具体例をいくつか紹介していきます。

パターン1: pandasのtransform()の挙動を再現

各groupの最大値や平均値などの統計量を新しい変数として追加したいときありますよ。そう言う時pandasではgroupby().transform()って書くだけでよかったのですが、polarsにはそれに直接対応する関数がありません。前まで仕方なくgroupbyした結果をjoin()していたのですが、over()を使えばいいことを知った時は感動しました。
書き方としては

df = df.with_columns(
    pl.col("最大値を出したい変数").max().over("グループ化したい変数").name.suffix("_max")
)

とすれば新しく変数ができます!

パターン2: グループごとでラグを取りたい

各グループごとでラグを取りたい時はめちゃめちゃ多いと思います。ラグを取る時はshift()を使えばいいのですが、さらにover使えばグループごとにラグ取れます。

df = df.with_columns(
    pl.col('value').shift().over('object')
)

めっちゃ楽ですね。

パターン3: グループごとにIDをつけたい

EDAするときなどにとりあえずグループごとにIDをつけて新しい変数として保存しておきたいなってことあると思います。その時はこう書きましょう（他の関数の説明は割愛させていただきます）。

df = df.with_columns(
        df.with_row_count('idx').select(pl.first('idx').over("グループ化したい変数")\
                                         .rank('dense') -1)
)

新しくidxという変数ができてそこに各グループごとのIDが格納されています。

他にもover関数使えばいろんなことできるので試してみてください！

〇〇の時、△△にして、それ以外の時××にする

この書き方はpolarsについて調べたら至る所に出てくると思いますが、個人的にめっちゃ好きな書き方なので載せておきます。

df = df.with_columns(
	pl.when(pl.col('x') >1)\
	.then(pl.lit(1))\
	.otherwise(pl.lit(0)).alias('x_over')
)

こんな簡単に条件分岐して新しい変数を作れるのは便利で好きです。

その他

便利な関数ではないのですが、こんな書き方もできるよと紹介したいものがあるので最後に補足として書いておきます。alias()を使わずに同じ挙動を再現できるんです。(上で使ってますが）変数の定義のような書き方をwith_columnsの中にすればその列が新しく生成されます。読みやすさ的にそちらの方が好きなのでこっちでいつもは書いています。
具体的に説明すると

df = df.with_columns(
    ~~~色々~~~.alias('新しい変数名')
)

という書き方と

df = df.with_columns(
    新しい変数名 = ~~~色々~~~
)

という書き方は同じ結果を返してくれます。この変数ってどう定義したっけ？ってあとでコード見返す時にこっちの方が圧倒的にみやすいと思うのでおすすめです。

最後に

もう少し紹介したい関数あるのですが、ニッチすぎるな...とかpandasに同じ処理する関数あるからな...というものばかりなので今回は控えさせていただきました。またもしストックが貯まれば第2弾など出していきたいと思います！

参考