Polars实现按日期分钟行转列并增量聚合（未来值置0）

需求：将按日期、分钟维度记录的价格数据转换为宽表，要求当前及之前分钟保留价格，之后分钟置0，全程采用向量化操作保证性能。

示例输入 #

import polars as pl

df = pl.DataFrame({
    "date": ["2022-01-01", "2022-01-01", "2022-01-02", "2022-01-02", 
             "2022-01-02", "2022-01-03", "2022-01-03", "2022-01-03"],
    "minute": [1, 2, 1, 2, 3, 1, 2, 3],
    "price": [10, 20, 15, 10, 20, 30, 60, 70]
})

解决方案代码 #

# 1. 按日期透视生成各分钟价格列
pivoted = df.pivot(
    index="date",
    columns="minute",
    values="price"
).rename(lambda col: f"{col}_price" if col != "date" else col)

# 2. 关联原表与透视表，让每条记录获取当日全量分钟价格
merged = df.join(pivoted, on="date")

# 3. 向量化判断并置0：当前分钟小于列对应分钟时设为0
minute_cols = [col for col in merged.columns if col.endswith("_price")]
min_numbers = [int(col.split("_")[0]) for col in minute_cols]

result = merged.with_columns(
    [
        pl.when(pl.col("minute") >= min_num)
        .then(pl.col(col))
        .otherwise(0)
        .alias(col)
        for col, min_num in zip(minute_cols, min_numbers)
    ]
)

# 还原原表顺序
result = result.sort(["date", "minute"])

print(result)

代码逻辑说明 #

• 透视宽表：通过pivot一次性聚合每个日期下的所有分钟价格，生成包含所有分钟列的中间表，这一步是Polars内置的向量化操作，效率极高
• 关联数据：将原始明细数据与透视表按date关联，让每条原始记录都能拿到当日所有分钟的价格值
• 向量化条件赋值：利用pl.when向量化判断逻辑，对每个分钟列批量处理：如果当前记录的minute小于该列对应的分钟数，就将值置0，否则保留原价格。完全避免Python循环，确保大数据场景下的性能
• 排序还原：最后按date和minute排序，保证输出顺序与原始输入一致

输出结果 #

shape: (8, 5)
┌────────────┬────────┬─────────┬─────────┬───────────┐
│ date       ┆ minute ┆ 1_price ┆ 2_price ┆ 3_price   │
│ ---        ┆ ---    ┆ ---     ┆ ---     ┆ ---       │
│ str        ┆ i64    ┆ i64     ┆ i64     ┆ i64       │
╞════════════╪════════╪═════════╪═════════╪═══════════╡
│ 2022-01-01 ┆ 1      ┆ 10      ┆ 0       ┆ 0         │
│ 2022-01-01 ┆ 2      ┆ 10      ┆ 20      ┆ 0         │
│ 2022-01-02 ┆ 1      ┆ 15      ┆ 0       ┆ 0         │
│ 2022-01-02 ┆ 2      ┆ 15      ┆ 10      ┆ 0         │
│ 2022-01-02 ┆ 3      ┆ 15      ┆ 10      ┆ 20        │
│ 2022-01-03 ┆ 1      ┆ 30      ┆ 0       ┆ 0         │
│ 2022-01-03 ┆ 2      ┆ 30      ┆ 60      ┆ 0         │
│ 2022-01-03 ┆ 3      ┆ 30      ┆ 60      ┆ 70        │
└────────────┴────────┴─────────┴─────────┴───────────┘

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