如何使用Polars更高效地按分组对列执行unstack操作？

需要在Polars中按列a分组，对b、c列执行unstack操作。当前通过循环分组处理的方式虽然能得到预期结果，但效率较低，寻求更高效的实现方式。

示例数据 #

import polars as pl

df = pl.from_repr("""
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ str │
╞═════╪═════╪═════╡
│ 1   ┆ a   ┆ f   │
│ 2   ┆ b   ┆ g   │
│ 1   ┆ c   ┆ h   │
│ 1   ┆ d   ┆ i   │
│ 2   ┆ e   ┆ j   │
└─────┴─────┴─────┘
""")

当前低效实现 #

seq_ = df["a"].unique()
l = len(seq_)
dfs = []

for idx, a in enumerate(seq_):
    df_tmp = df.filter(pl.col("a") == a)
    steps = df_tmp.shape[0]
    df_tmp = df_tmp.unstack(step=steps, how="horizontal", columns=["b", "c"])
    df_tmp = df_tmp.with_columns(pl.lit(a).alias("a"))
    dfs.append(df_tmp)

df_tmp = pl.concat(dfs, how="diagonal")

预期输出 #

┌─────┬─────┬─────┬──────┬─────┬─────┬──────┐
│ a   ┆ b_0 ┆ b_1 ┆ b_2  ┆ c_0 ┆ c_1 ┆ c_2  │
│ --- ┆ --- ┆ --- ┆ ---  ┆ --- ┆ --- ┆ ---  │
│ i64 ┆ str ┆ str ┆ str  ┆ str ┆ str ┆ str  │
╞═════╪═════╪═════╪══════╪═════╪═════╪══════╡
│ 1   ┆ a   ┆ c   ┆ d    ┆ f   ┆ h   ┆ i    │
│ 2   ┆ b   ┆ e   ┆ null ┆ g   ┆ j   ┆ null │
└─────┴─────┴─────┴──────┴─────┴─────┴──────┘

高效实现方式 #

可以利用Polars的分组聚合+列表转结构体+展开的向量化操作，避免循环带来的性能损耗：

result = (
    df
    .group_by("a", maintain_order=True)
    .agg(
        pl.col("b").list.to_struct(n_field_strategy="max_width"),
        pl.col("c").list.to_struct(n_field_strategy="max_width")
    )
    .unnest("b", "c")
)

print(result)

代码说明 #

1. group_by("a")：按列a分组，maintain_order=True保留分组的原始顺序
2. agg(...)：对每组的b、c列分别聚合为列表，再通过list.to_struct将列表转换为结构体，n_field_strategy="max_width"会自动以所有分组中最长的列表长度作为结构体的字段数，不足的位置填充null
3. unnest("b", "c")：将结构体字段展开为单独的列，自动生成b_0、b_1、c_0这类命名的列

这种方式完全是向量化操作，性能远优于循环处理，尤其在数据量较大时提升明显。

内容的提问来源于stack exchange，提问作者wedrano de carvalho