如何将Polars DataFrame转换为以日期为列名的格式？

问题描述 #

我正计划把工作里的部分Pandas代码迁移到Polars。现有Pandas DataFrame是宽格式，每行对应用户每月的金额数据：

2023-09-01  2023-10-01  ...  2024-09-01  2024-10-01
customer_id                                                                                                                                                                         
111111          0.000000    0.000000  ...  918.333333  918.333333
222222        466.666667  883.333333  ...    0.000000    0.000000
333333          0.000000  833.333333  ...  833.333333    0.000000
444444          0.000000    0.000000  ...  833.333333  833.333333

现在我有一个Polars的长格式DataFrame，结构如下：

shape: (313_590, 3)
┌─────────────┬─────────────────────┬────────────┐
│ customer_id ┆ date                ┆ amount     │
│ ---         ┆ ---                 ┆ ---        │
│ i64         ┆ datetime[μs]        ┆ f64        │
╞═════════════╪═════════════════════╪════════════╡
│ 111111      ┆ 2023-01-01 00:00:00 ┆ 80.749008  │
│ 222222      ┆ 2023-06-01 00:00:00 ┆ 87.628968  │
│ 333333      ┆ 2023-02-01 00:00:00 ┆ 180.327381 │
│ 333333      ┆ 2022-06-01 00:00:00 ┆ 180.327381 │
│ …           ┆ …                   ┆ …          │
│ 555555      ┆ 2022-05-01 00:00:00 ┆ 85.818452  │
│ 666666      ┆ 2022-06-01 00:00:00 ┆ 85.818452  │
│ 777777      ┆ 2023-11-01 00:00:00 ┆ 87.628968  │
│ 888888      ┆ 2023-12-01 00:00:00 ┆ 87.628968  │
└─────────────┴─────────────────────┴────────────┘

比如客户111111可能有12行数据对应每个月的记录。请问怎么在Polars里把这个DataFrame转换成上面的宽格式，让列名是日期的ISO格式字符串？

解决方案 #

可以通过Polars的pivot方法结合日期格式转换实现需求，具体步骤如下：

1. 将日期列转换为ISO格式字符串 #

先把datetime类型的date列转换成YYYY-MM-DD格式的字符串，作为宽格式的列名：

import polars as pl

df = df.with_columns(
    pl.col("date").dt.strftime("%Y-%m-%d").alias("date_str")
)

2. 执行pivot转换为宽格式 #

使用pivot方法，以customer_id作为行标识，date_str作为列名，amount作为值，并将缺失的月份填充为0：

wide_df = df.pivot(
    index="customer_id",
    columns="date_str",
    values="amount",
    aggregate_function="first"  # 每个客户每月仅一条记录，用first即可
).fill_null(0.0)

完整示例代码 #

import polars as pl

# 模拟原始长格式DataFrame
df = pl.DataFrame({
    "customer_id": [111111, 222222, 333333, 333333],
    "date": [
        pl.datetime(2023, 1, 1),
        pl.datetime(2023, 6, 1),
        pl.datetime(2023, 2, 1),
        pl.datetime(2022, 6, 1)
    ],
    "amount": [80.749008, 87.628968, 180.327381, 180.327381]
})

# 转换日期格式
df = df.with_columns(
    pl.col("date").dt.strftime("%Y-%m-%d").alias("date_str")
)

# 生成宽格式DataFrame
wide_df = df.pivot(
    index="customer_id",
    columns="date_str",
    values="amount",
    aggregate_function="first"
).fill_null(0.0)

print(wide_df)

关键说明 #

• dt.strftime("%Y-%m-%d")：将datetime类型转为标准ISO日期字符串，确保列名格式与目标一致。
• aggregate_function：如果同一客户同一月份存在多条记录，可替换为sum、mean等聚合函数；单条记录时用first即可。
• fill_null(0.0)：填充客户无记录的月份为0，与你提供的Pandas示例格式匹配。

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