如何在Polars中基于单列分组创建新列（宽表转换）

问题背景 #

初始长表结构：

期望转换为宽表，以时间为行索引，每个id作为单独列：

已通过Pandas实现该功能，但因数据量庞大且需多次调用，希望改用Polars提升处理速度，尝试过groupby及join/hstack/concat方法，但在使用LazyFrame时遇到问题。

附数据生成代码：

import polars as pl

df = pl.DataFrame({'id': [2050, 2051, 2050, 2051, 2051],
                    'time': ['2023-05-01',
                             '2023-05-01',
                             '2023-05-02',
                             '2023-05-02',
                             '2023-05-03'],
                   'value': [20, 25, 21, 22, 23]})
df = df.with_columns(pl.col("time").str.to_datetime("%Y-%m-%d"))

解决方案 #

Polars内置的pivot方法专门解决这类长表转宽表的需求，不管是DataFrame还是LazyFrame都能高效处理，无需复杂的分组或拼接操作。

1. 处理普通DataFrame #

直接调用pivot，指定行索引、列字段和数值列即可：

wide_df = df.pivot(
    index="time",
    columns="id",
    values="value"
)
print(wide_df)

输出结果：

shape: (3, 3)
┌─────────────────────┬──────┬──────┐
│ time                ┆ 2050 ┆ 2051 │
│ ---                 ┆ ---  ┆ ---  │
│ datetime[μs]        ┆ i64  ┆ i64  │
╞═════════════════════╪══════╪══════╡
│ 2023-05-01 00:00:00 ┆ 20   ┆ 25   │
│ 2023-05-02 00:00:00 ┆ 21   ┆ 22   │
│ 2023-05-03 00:00:00 ┆ null ┆ 23   │
└─────────────────────┴──────┴──────┘

2. 处理LazyFrame（大数据场景） #

如果用LazyFrame优化内存使用，只需将pivot加入懒加载流水线，最后执行collect()即可：

lf = df.lazy()
wide_lf = lf.pivot(
    index="time",
    columns="id",
    values="value"
)
wide_df = wide_lf.collect()
print(wide_df)

输出结果和上面一致，全程懒加载，适合处理超大数据集。

自定义缺失值填充 #

如果需要把默认的null替换成指定值（比如0），可以添加aggregate_function参数：

wide_df = df.pivot(
    index="time",
    columns="id",
    values="value",
    aggregate_function=lambda x: x.first().fill_null(0)
)

内容的提问来源于stack exchange，提问作者The-Viggo-Meister