如何在Polars中按ID和45分钟时间阈值分组生成新列？

需求说明 #

现有包含col1（分组键）和col2（已转换为datetime类型）的Polars DataFrame，需要按col1分组，将每组内的col2按45分钟时间区间划分，生成col3列——同一45分钟区间内的记录col3值相同，后续区间值依次递增。

示例数据 #

import polars as pl

data = {
    "col1": [1, 1, 1,1,1,1,1,1,1,1,1,1, 2, 2,2,2,2,2,2,2,2,2,2],
    "col2": [
        "2022-05-25T08:00:00.648681",
        "2022-05-25T08:15:00.648681",
        "2022-05-25T08:30:00.648681",
        "2022-05-25T08:45:00.648681",
        "2022-05-25T09:00:00.648681",
        "2022-05-25T09:15:00.648681",
        "2022-05-25T09:30:00.648681",
        "2022-05-25T09:45:00.648681",
        "2022-05-25T10:00:00.648681",
        "2022-05-25T10:15:00.648681",
        "2022-05-25T10:30:00.648681",
        "2022-05-25T10:45:00.648681",
        "2022-05-25T08:00:00.648681",
        "2022-05-25T08:15:00.648681",
        "2022-05-25T08:30:00.648681",
        "2022-05-25T08:45:00.648681",
        "2022-05-25T09:00:00.648681",
        "2022-05-25T06:00:00.648681",
        "2022-05-25T06:15:00.648681",
        "2022-05-25T06:30:00.648681",
        "2022-05-25T06:45:00.648681",
        "2022-05-25T07:00:00.648681",
        "2022-05-25T07:15:00.648681",
    ],
}

df = pl.DataFrame(data)
df = df.with_columns(pl.col("col2").str.to_datetime())

期望输出 #

┌──────┬────────────────────────────┬──────┐
│ col1 ┆ col2                       ┆ col3 │
│ ---  ┆ ---                        ┆ ---  │
│ i64  ┆ datetime[μs]               ┆ u32  │
╞══════╪════════════════════════════╪══════╡
│ 1    ┆ 2022-05-25 08:00:00.648681 ┆ 1    │
│ 1    ┆ 2022-05-25 08:15:00.648681 ┆ 1    │
│ 1    ┆ 2022-05-25 08:30:00.648681 ┆ 1    │
│ 1    ┆ 2022-05-25 08:45:00.648681 ┆ 1    │
│ 1    ┆ 2022-05-25 09:00:00.648681 ┆ 2    │
│ 1    ┆ 2022-05-25 09:15:00.648681 ┆ 2    │
│ 1    ┆ 2022-05-25 09:30:00.648681 ┆ 2    │
│ 1    ┆ 2022-05-25 09:45:00.648681 ┆ 2    │
│ 1    ┆ 2022-05-25 10:00:00.648681 ┆ 3    │
│ 1    ┆ 2022-05-25 10:15:00.648681 ┆ 3    │
│ 1    ┆ 2022-05-25 10:30:00.648681 ┆ 3    │
│ 1    ┆ 2022-05-25 10:45:00.648681 ┆ 3    │
│ 2    ┆ 2022-05-25 08:00:00.648681 ┆ 3    │
│ 2    ┆ 2022-05-25 08:15:00.648681 ┆ 3    │
│ 2    ┆ 2022-05-25 08:30:00.648681 ┆ 3    │
│ 2    ┆ 2022-05-25 08:45:00.648681 ┆ 3    │
│ 2    ┆ 2022-05-25 09:00:00.648681 ┆ 4    │
│ 2    ┆ 2022-05-25 06:00:00.648681 ┆ 1    │
│ 2    ┆ 2022-05-25 06:15:00.648681 ┆ 1    │
│ 2    ┆ 2022-05-25 06:30:00.648681 ┆ 1    │
│ 2    ┆ 2022-05-25 06:45:00.648681 ┆ 1    │
│ 2    ┆ 2022-05-25 07:00:00.648681 ┆ 2    │
│ 2    ┆ 2022-05-25 07:15:00.648681 ┆ 2    │
└──────┴────────────────────────────┴──────┘

解决方案 #

方法一：时间向下取整+稠密排名 #

通过将时间戳向下取整到45分钟区间，再对分组内的取整结果做稠密排名，实现区间标记：

result_df = df.group_by("col1", maintain_order=True).agg(
    pl.col("col2"),
    col3=pl.col("col2")
        .dt.floor("45m")  # 将时间向下取整到最近的45分钟起始点
        .rank(method="dense", order="asc")  # 对取整结果做稠密排名，同一区间排名相同
).explode(["col2", "col3"])  # 将聚合后的列表展开回原行结构

print(result_df)

方法二：计算时间差+区间编号 #

通过计算每个时间点与分组内最早时间的差值，转换为分钟数后整除45，得到区间编号：

result_df = df.group_by("col1", maintain_order=True).agg(
    pl.col("col2"),
    col3=(pl.col("col2") - pl.col("col2").min())  # 计算与分组内最早时间的差值
        .dt.total_minutes()  # 转换为总分钟数
        .floordiv(45) + 1  # 整除45得到区间索引（从0开始），加1让编号从1开始
).explode(["col2", "col3"])

print(result_df)

两种方法都能生成符合要求的col3列，可根据习惯选择。

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