Polars计算：连续标签分组内与组首元素的时间差

要实现需求，核心是先把连续出现的相同group标记为独立分组，再基于这些分组计算每条数据与分组内首个timestamp的时间差。

完整代码实现 #

import polars as pl

df = pl.from_repr("""
┌─────────────────────────┬─────────┐
│ timestamp               ┆ group   │
│ ---                     ┆ ---     │
│ datetime[μs, UTC]       ┆ str     │
╞═════════════════════════╪═════════╡
│ 2009-04-18 11:30:00 UTC ┆ group_1 │
│ 2009-04-18 11:40:00 UTC ┆ group_1 │
│ 2009-04-18 11:50:00 UTC ┆ group_1 │
│ 2009-04-18 12:00:00 UTC ┆ group_2 │
│ 2009-04-18 12:10:00 UTC ┆ group_2 │
│ 2009-04-18 12:20:00 UTC ┆ group_1 │ # <- reappearance of group_1
│ 2009-04-18 12:30:00 UTC ┆ group_1 │ # <- reappearance of group_1
└─────────────────────────┴─────────┘
""")

# 生成连续分组标识并计算时间差
result = df.with_columns(
    # 标记连续分组：当前group与上一行不同时触发新分组，累加得到分组ID
    group_id=pl.col("group").ne(pl.col("group").shift()).cumsum()
).group_by("group_id", maintain_order=True).agg(
    pl.all(),
    # 计算每个分组内timestamp与首条的时间差，转换为分钟整数
    timediff=(pl.col("timestamp") - pl.col("timestamp").first()).dt.total_minutes().cast(pl.Int32)
).explode(pl.exclude("group_id")).drop("group_id")

print(result)

代码说明 #

1. 生成连续分组ID：

   • 使用pl.col("group").ne(pl.col("group").shift())判断当前行的group是否与上一行不同，不同则返回True（1），相同返回False（0）。
   • 再用cumsum()累加这些布尔值，得到唯一的group_id，确保连续相同的group属于同一个分组，后续再次出现的group会被标记为新分组。

2. 计算时间差：

   • 按group_id分组时开启maintain_order=True，保证结果顺序和原数据一致。
   • 用pl.col("timestamp") - pl.col("timestamp").first()计算每条数据与分组内首个时间的差值，再通过dt.total_minutes()转换为分钟数，最后转成整数类型匹配示例格式。

3. 整理结果：

   • 用explode展开分组聚合后的结果，删除临时的group_id列，得到最终格式。

输出结果 #

┌─────────────────────────┬─────────┬──────────┐
│ timestamp               ┆ group   ┆ timediff │
│ ---                     ┆ ---     ┆ ---      │
│ datetime[μs, UTC]       ┆ str     ┆ i32      │
╞═════════════════════════╪═════════╪══════════╡
│ 2009-04-18 11:30:00 UTC ┆ group_1 ┆ 0        │
│ 2009-04-18 11:40:00 UTC ┆ group_1 ┆ 10       │
│ 2009-04-18 11:50:00 UTC ┆ group_1 ┆ 20       │
│ 2009-04-18 12:00:00 UTC ┆ group_2 ┆ 0        │
│ 2009-04-18 12:10:00 UTC ┆ group_2 ┆ 10       │
│ 2009-04-18 12:20:00 UTC ┆ group_1 ┆ 0        │
│ 2009-04-18 12:30:00 UTC ┆ group_1 ┆ 10       │
└─────────────────────────┴─────────┴──────────┘

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