集約統計量はPolarsが作りやすいぞ

自己紹介

データサイエンティストとして仕事をしているeikichiです。
Qiita初投稿です。

背景

KaggleのOTTOコンペ参加中にpolarsを知りました。
その時はメモリも圧縮されてさらに処理が早いというところで使用していましたが
処理を重ねて書いていけるので実装しやすいし、率先して使うようになっていきました。
今回はPolarsの紹介として集約特徴量の作りやすさについて紹介します。

対象者

• Polars初心者
• pandasとpolarsの書き方の違いをざっくり把握したい

概要

今回は特徴量エンジニアリングをするうえでほぼ作成する集約統計量をどのようにかけるかを前半はpandasと比較も入れながら説明していきます。

環境: kaggle notebook
polars: 2.0.3
panas: 0.19.12

データ

Kaggle Predict Student Performance from Game Playのtrainデータ(4.72GB)を使用します。

本コンペのtrainデータではsession_id毎のゲームプレイの時系列データが与えられ、推論ではゲーム内での3つのチェックポイントで出題される問題に対して各プレイヤーが正しく答えられているかを予測する課題でした。

データを様々な条件で集約統計量にして特徴量を作成することが多かったのでこちらを選定しました。

実験

各々のデータフレームで読み込み。
読み込みでもpolarsのDataFrameの方が1/2の速度で完了してる。速い。

まず、基本操作としてgroupbyを使って”session_id”毎のカウントをとることにします。

Pandas

(
    pd_df
    .groupby("session_id")
    .agg(count=("index", "count"))
    .sort_values("count")
    .reset_index()
    .head()
)

Polars

(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(pl.col("index").count().alias("count"))
    .sort("count")
    .head()
)

Pandasではaggの引数に{rename後変数名}=(”{集約変数}”, “{集約関数}”)を渡すことで指定した集約変数を集約関数で集計後、カラム名をrenameしてデータフレームに返すことができます。

Polarsではgroup_byメソッドで集計keyを渡します。
aggメソッドではpl.col({”集約変数}).count()で集約し、aliasメソッドにrename後の変数名を渡してカラム名を変更できます。

Polarsの特徴として、polars.Expressionを使います。pl.colのやつです。
Pandasライクにdf[”index”]のような指定もできますが、深く使おうとしていくとつまづきポイントになります。（具体的にはlazyframe使いたいときや、後述など）なのでPolars使うならExpression表記に慣れていくほうがいいでしょう。

次に”session_id”毎の”elapsed_time”の合計をとることにします。

Pandas

%%time
(
    pd_df
    .groupby("session_id")
    .agg(
        index_count=("index", "count"),
        elasped_time_sum=("elapsed_time", "sum")
    )
    .sort_values("index_count")
    .reset_index()
    .head()
)

Polars

%%time
(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(
        pl.col("index").count().alias("index_count"),
        pl.col("elapsed_time").sum().alias("elapsed_time_sum"),
    )
    .sort("index_count")
    .head()
)

両ともに先ほど説明した構文で１行実装を追加して、elapsed_time合計列を追加することができました。
ここで、aggの中身をリストで渡すとどのようになるのでしょうか。

Pandas

Polars

Pandasではそもそも構文をリストにするときにエラーが出てしまい、置き換えることができませんがPolarsでは問題なく出力ができています。

急に引数をリストで渡したりして「なんで？」となっているかと思いますが、リストで渡せることによって以下のように書けていきます。

Polars

(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(
        [
            pl.col("index").count().alias("index_count"),
            *[pl.col(c).sum().alias(f"{c}_sum") for c in ["elapsed_time", "index", "fullscreen", "hq", "music"]],
        ]
    )
    .sort("index_count")
    .head()
)

aggの2行目をリスト内包表記を使い、合計の集約特徴量を算出したいカラムをfor文で回してpolars.Expression表記を作成しています。
具体的には、上記は以下の実装と同等のものになります。

Polars

(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(
        [
            pl.col("index").count().alias("index_count"),
            pl.col("elapsed_time").sum().alias("elapsed_time_sum"),
            pl.col("index").sum().alias("index_sum"),
            pl.col("fullscreen").sum().alias("fullscreen_sum"),
            pl.col("hq").sum().alias("hq_sum"),
            pl.col("music").sum().alias("music_sum"),
        ]
    )
    .sort("index_count")
    .head()
)

polars.Expression表記を柔軟に使うことで上記の実装を1行でさくっと集約特徴量を作れるので便利ですよね。

最初にリストで渡せるというところを確認しましたが、その部分を生かして実行部分と分割して実装することができます。（集約特徴量も増やしてみました。）

Polars

# 集約したい変数リスト
NUMS = ["elapsed_time", "index", "fullscreen", "hq", "music"]
# aggに渡す引数
aggs = [
            pl.col("index").count().alias("index_count"),
            *[pl.col(c).sum().alias(f"{c}_sum") for c in NUMS],
            *[pl.col(c).mean().alias(f"{c}_mean") for c in NUMS], 
            *[pl.col(c).std().alias(f"{c}_std") for c in NUMS],
]

# 実行処理部分
(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(aggs)
    .sort("index_count")
    .head()
)

さらに、Polarsではあるカラムの状態を条件において集約特徴量を作ることもaggの中でできます。
具体的にはfilterメソッドを追加して集約特徴量を作るようにします。

今回は”event_name”の水準に応じての”session_id”毎のカウントをとるようにします。

Polars

データフレーム表示が長くなったので一部集計特徴量をコメントアウトしています。
追記したのはaggsのリスト最後行です。filterメソッドを条件を加え、リスト内包表記で全体のpolars.Expression表記を作成しています。

さらに応用して、リスト内包表記を重ねることであるカラムの状態を条件における指定カラムごとの集約特徴量も以下のように作成することができます。

Polars

# フィルタリスト
EVENT_NAMES = pl_df.get_column("event_name").unique().to_list()
# 集約したい変数リスト
NUMS = ["elapsed_time", "index", "fullscreen", "hq", "music"]
# aggに渡す引数
aggs = [
            pl.col("index").count().alias("index_count"),
            *[pl.col(c).sum().alias(f"{c}_sum") for c in NUMS],
            *[pl.col(c).mean().alias(f"{c}_mean") for c in NUMS], 
            *[pl.col(c).std().alias(f"{c}_std") for c in NUMS],
            *[pl.col("index").filter(pl.col("event_name")==v).count().alias(f"{v}_count") for v in EVENT_NAMES],
            *[pl.col(c).filter(pl.col("event_name")==v).sum().alias(f"{c}_{v}_sum") for c in NUMS for v in EVENT_NAMES],
]

# 実行処理部分
(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(aggs)
    .sort("index_count")
    .head()
)

最後に上記の内容をすべて反映させて特徴量を作成すると、aggに8行ほど渡すだけで200個ほどの集約特徴量を作成することができました。

Polars

# フィルタリスト
EVENT_NAMES = pl_df.get_column("event_name").unique().to_list()
# 集約したい変数リスト
NUMS = ["elapsed_time", "index", "fullscreen", "hq", "music"]
# aggに渡す引数
aggs = [
            pl.col("index").count().alias("index_count"),
            *[pl.col(c).sum().alias(f"{c}_sum") for c in NUMS],
            *[pl.col(c).mean().alias(f"{c}_mean") for c in NUMS], 
            *[pl.col(c).std().alias(f"{c}_std") for c in NUMS],
            *[pl.col("index").filter(pl.col("event_name")==v).count().alias(f"{v}_count") for v in EVENT_NAMES],
            *[pl.col(c).filter(pl.col("event_name")==v).sum().alias(f"{c}_{v}_sum") for c in NUMS for v in EVENT_NAMES],
            *[pl.col(c).filter(pl.col("event_name")==v).mean().alias(f"{c}_{v}_mean") for c in NUMS for v in EVENT_NAMES],
            *[pl.col(c).filter(pl.col("event_name")==v).std().alias(f"{c}_{v}_std") for c in NUMS for v in EVENT_NAMES],
]

# 実行処理部分
(
    pl_df
    .group_by("session_id", maintain_order=True)
    .agg(aggs)
    .sort("index_count")
    .head()
)

まとめ

Polarsでの集約特徴量の作りやすさについて紹介しました。

filterの中身もbool値が返れば良いので.is_inなど様々な条件を渡すことで意図する集計特徴量が作れます。今回はgroup_byにフォーカス当てましたが、with_columnsメソッドでも同じことが言えます。
なので強みはpolars.Expression表記が処理内で柔軟に使えるところかと解釈しています。

実際には、増やすことが重要ではなく、意味がある特徴量を作成するのが本質なので検証しながら進めていくことになります。ただ、作りやすいのはメリットだし表記の仕方を統一感を持って実装できるので良きです。あとやっぱり速い。