基于Python Polars库实现随Col1值变化修改Col2日期月份

要实现当Col1值变化时更新Col2的月份，可通过以下步骤完成：

步骤说明 #

1. 识别Col1的变化点：判断当前行Col1是否与前一行不同，标记变化位置。
2. 计算月份增量：对变化点进行累计求和，得到每个分组对应的月份编号。
3. 更新日期月份：利用Polars的日期处理函数，将Col2的月份按编号调整，Polars会自动处理无效日期（如31号转2月会调整为当月最后一天）。

代码实现 #

from datetime import date
import polars as pl

my_df = pl.DataFrame({
    'Col1': [1, 1, 1, 1, 2, 2, 3, 3],
    'Col2': [date(2020, 1, 31)] * 8,
})

# 处理逻辑
result_df = my_df.with_columns(
    # 标记Col1的变化行（第一行默认标记为变化）
    change_flag=pl.col("Col1").ne(pl.col("Col1").shift(1)).fill_null(True),
    # 累计变化次数，作为目标月份的编号
    target_month=pl.col("change_flag").cumsum()
).with_columns(
    # 更新Col2的月份
    Col2=pl.col("Col2").dt.set_month(pl.col("target_month"))
).drop(["change_flag", "target_month"])

print(result_df)

输出结果 #

shape: (8, 2)
┌──────┬────────────┐
│ Col1 ┆ Col2       │
│ ---  ┆ ---        │
│ i64  ┆ date       │
╞══════╪════════════╡
│ 1    ┆ 2020-01-31 │
│ 1    ┆ 2020-01-31 │
│ 1    ┆ 2020-01-31 │
│ 1    ┆ 2020-01-31 │
│ 2    ┆ 2020-02-29 │
│ 2    ┆ 2020-02-29 │
│ 3    ┆ 2020-03-31 │
│ 3    ┆ 2020-03-31 │
└──────┴────────────┘

扩展说明 #

如果希望基于Col2的原始月份递增（而非固定从1月开始），可改用dt.offset_by函数实现：

result_df = my_df.with_columns(
    change_flag=pl.col("Col1").ne(pl.col("Col1").shift(1)).fill_null(True),
    # 计算需要增加的月份数（第一次变化加0，第二次加1，以此类推）
    month_add=pl.col("change_flag").cumsum() - 1
).with_columns(
    Col2=pl.col("Col2").dt.offset_by(pl.col("month_add").cast(str) + "mo")
).drop(["change_flag", "month_add"])

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