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

# クエリ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()

# クエリ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"]]

# クエリ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()

# 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 = 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"]]

pandas==1.5.2
polars==0.14.31

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)
)