如何简化Polars DataFrame分组透视并生成自定义列名的实现？

需求说明 #

需要对Polars DataFrame按a列分组，将b列与c列组合成字符串标识符，以此作为透视后DataFrame的列名。现有代码可实现需求，但实现路径繁琐，希望得到更简洁的写法（数据规模最大1000行，无需考虑性能优化）。

示例数据与原始代码 #

import polars as pl

data = {
    "a": [1, 1, 1, 2, 2, 3],
    "b": [11, 12, 13, 11, 12, 11],
    "c": ["x1", "x2", "x3", "x1", "x2", "x1"],
    "val": [101, 102, 102, 201, 202, 301],
}
df = pl.DataFrame(data)

原始实现代码：

counter = 0
for tmp_df in df.partition_by("a"):
    grp_df = (
        tmp_df.with_columns((pl.col("b").cast(pl.String) + "_" + pl.col("c")).alias("col_id"))
        .drop("b", "c")
        .pivot("col_id", index="a")
    )

    if counter == 0:
        result_df = grp_df.select(pl.all())
    else:
        result_df = pl.concat([result_df, grp_df], how="diagonal")
    counter += 1

print(result_df)

原始输出：

shape: (3, 4)
┌─────┬───────┬───────┬───────┐
│ a   ┆ 11_x1 ┆ 12_x2 ┆ 13_x3 │
│ --- ┆ ---   ┆ ---   ┆ ---   │
│ i64 ┆ i64   ┆ i64   ┆ i64   │
╞═════╪═══════╪═══════╪═══════╡
│ 1   ┆ 101   ┆ 102   ┆ 102   │
│ 2   ┆ 201   ┆ 202   ┆ null  │
│ 3   ┆ 301   ┆ null  ┆ null  │
└─────┴───────┴───────┴───────┘

更简洁的实现方式 #

无需循环分区拼接，直接生成组合列后一次性完成透视即可：

result_df = (
    df
    .with_columns(col_id=pl.col("b").cast(pl.String) + "_" + pl.col("c"))
    .pivot(index="a", columns="col_id", values="val", aggregate_function="first")
)

print(result_df)

代码解释 #

1. 生成组合列：通过with_columns将b转为字符串后与c拼接成col_id，作为透视后的列名来源。
2. 一次性透视：调用pivot方法，指定index="a"作为分组依据，columns="col_id"将组合列转为透视后的列，values="val"填充对应单元格的值，aggregate_function="first"确保每个分组下的唯一对应值被保留（因数据中每个a+col_id组合仅对应一个val，用first/last/sum均可）。

输出结果 #

与原始代码输出完全一致：

shape: (3, 4)
┌─────┬───────┬───────┬───────┐
│ a   ┆ 11_x1 ┆ 12_x2 ┆ 13_x3 │
│ --- ┆ ---   ┆ ---   ┆ ---   │
│ i64 ┆ i64   ┆ i64   ┆ i64   │
╞═════╪═══════╪═══════╪═══════╡
│ 1   ┆ 101   ┆ 102   ┆ 102   │
│ 2   ┆ 201   ┆ 202   ┆ null  │
│ 3   ┆ 301   ┆ null  ┆ null  │
└─────┴───────┴───────┴───────┘

内容的提问来源于stack exchange，提问作者darked89