Zenn
Closed10

pandasとpolarsのクエリ実行時間の比較

bilzardbilzard

pandasとpolarsのクエリ実行時間の比較。
以下のクエリを比較した。
同等の処理内容のクエリをpandasで書けなかったので、測定結果にはクエリの書き方による差異も含まれている。

pandasのクエリ

# クエリ1
pd.concat(
    [
        train_pandas.groupby("session")["aid"].transform(lambda x: x),
        train_pandas.groupby("session")["aid"]
        .transform(lambda x: x.shift(-1))
        .rename("aid_next"),
    ],
    axis=1,
).dropna()

polarsのクエリ

# クエリ2
train_pairs = (
    pl.concat([train, test])
    .groupby("session")
    .agg([pl.col("aid"), pl.col("aid").shift(-1).alias("aid_next")])
    .explode(["aid", "aid_next"])
    .drop_nulls()
)[["aid", "aid_next"]]

クエリ実行時間の比較

N pandas polars pandas/polars
2000000 1.95 0.0641 30.421217
20000000 27.00 0.3910 69.053708
216716096 720.00 7.5600 95.238095

bilzardbilzard

ちなみにpandasでグループごとに列を結合するクエリは遅かった（おそらくconcatのオーバーヘッドが大きい）。

  • データ数: 2M
  • クエリ1: 2.69sec
  • クエリ3: 13.8sec
# クエリ3
def get_consecutive_aids(x):
    return pd.concat([x["aid"], x["aid"].shift(-1).rename("aid_next")], axis=1)


train_pandas.groupby("session").apply(get_consecutive_aids).reset_index().drop(
    "level_1", axis=1
).dropna()
bilzardbilzard

[追記] mac環境で最新のpandasで、クエリも修正して実行し直したらそこまで致命的な差ではなかった。
pandas==1.5.1, polars==0.14.31
最新のpandasはマルチプロセスに対応しているのか？
それでもpolarsが倍早い。

pyarrowのバージョンとかでも変わるのだろうか？

# pandas==1.5.1
pd.concat(
    [pd.concat([df["b"], df.groupby("a")["b"].shift(-1).rename("b_next")], axis=1)],
    axis=1,
).dropna()

# polars==0.14.31
pl_df.groupby("a").agg([pl.col("b"), pl.col("b").shift(-1).alias("b_next")]).explode(
    ["b", "b_next"]
).drop_nulls()[["b", "b_next"]]
N pandas polars pandas/polars
20000000 0.618 0.354 1.745763
50000000 1.680 0.765 2.196078
100000000 5.020 2.220 2.261261
200000000 11.200 4.970 2.253521

bilzardbilzard

pandasのクエリが冗長だったので測定し直した。
データの生成方法を変えたためか、polarsが上の例より速くなっている。
pandasとpolarsの速度の比は2-4倍。

N = 200_000_000
df = (
    pd.DataFrame(
        {
            "a": np.random.randint(0, N // 500, N).astype("uint32"),
            "b": np.random.randint(0, N // 100, N).astype("uint32"),
        }
    )
    .sort_values("a")
    .reset_index(drop=True)
)
df

%%time
pd.concat([df["b"], df.groupby("a")["b"].shift(-1).rename("b_next")], axis=1).dropna()

%%time
pl_df.groupby("a").agg([pl.col("b"), pl.col("b").shift(-1).alias("b_next")]).explode(
    ["b", "b_next"]
).drop_nulls()[["b", "b_next"]]
N pandas polars pandas/polars
20000000 0.694 0.390 1.779487
50000000 1.650 0.415 3.975904
100000000 4.100 1.380 2.971014
200000000 10.800 3.590 3.008357

bilzardbilzard

Ubuntu環境でもバージョンを揃えて同じことをしてみた。
クエリとデータの生成方法はmac環境のもの[1]と同じ。
結果はmac環境とほぼ同様。計算速度はほぼ線形で、polarsとpandasの速度比は概ね3倍。

pandas==1.5.2
polars==0.14.31
N pandas polars pandas/polars
20000000 0.547 0.202 2.707921
50000000 1.380 0.463 2.980562
100000000 2.770 0.948 2.921941
200000000 5.500 1.930 2.849741

脚注

  1. https://zenn.dev/link/comments/f46352dc2302cc ↩︎

bilzardbilzard

データのカーディナりてぃを変えて測定。
この条件ではpandasとpolarsがほぼ互角だった。

N = 50_000_000
df = (
    pd.DataFrame(
        {
            "a": np.random.randint(0, N // 20, N).astype("uint32"),
            "b": np.random.randint(0, N // 50, N).astype("uint32"),
        }
    )
    .sort_values("a")
    .reset_index(drop=True)
)
N pandas polars pandas/polars
20000000 0.288 0.263 1.095057
50000000 1.510 1.410 1.070922
100000000 2.990 2.770 1.079422
200000000 5.960 5.780 1.031142

bilzardbilzard

グループ数による実行速度の差異は、クエリの違いが原因かもしれない。

bilzardbilzard

結局一番最初の計測はpandasのクエリが悪すぎたよう。

このスクラップは2022/11/23にクローズされました