如何在Polars中按周与时段分组计算唯一用户数量？

现有一个包含3个字段的Polars DataFrame：

df = pl.from_repr("""
┌─────────┬─────────────────────┬─────────────┐
│ user_id ┆ date                ┆ part_of_day │
│ ---     ┆ ---                 ┆ ---         │
│ i64     ┆ datetime[ns]        ┆ cat         │
╞═════════╪═════════════════════╪═════════════╡
│ 173367  ┆ 2021-08-03 00:00:00 ┆ day         │
│ 132702  ┆ 2021-10-28 00:00:00 ┆ evening     │
│ 100853  ┆ 2021-07-29 00:00:00 ┆ night       │
│ 305810  ┆ 2021-08-24 00:00:00 ┆ day         │
│ 305239  ┆ 2021-08-13 00:00:00 ┆ day         │
└─────────┴─────────────────────┴─────────────┘
""")

任务是计算每周及各时段的唯一用户数量。已知Polars有三种分组方式：

• group_by：支持按列或表达式进行简单分组
• group_by_dynamic：仅接受datetime列作为分组键
• rolling：同样仅接受datetime列作为分组键

group_by_dynamic看似最适配需求，但无法同时将其他列作为分组键，请问如何用Polars完成该任务？

附Pandas实现代码：

(
    df
    .cast({"part_of_day": pl.String})
    .to_pandas()
    .groupby([pd.Grouper(key="date", freq="1W"), "part_of_day"])
    .user_id
    .nunique()
)

期望输出结果：

date        part_of_day
2021-08-01  night          1
2021-08-08  day            1
2021-08-15  day            1
2021-08-29  day            1
2021-10-31  evening        1
Name: user_id, dtype: int64

直接使用group_by结合日期截断表达式，同时按周截断后的日期和part_of_day分组，即可实现需求，无需依赖group_by_dynamic。

方法1：使用dt.truncate截断日期到周 #

这是最直接高效的方式：

result = (
    df
    .group_by(
        # 将日期截断到当周起始日，与Pandas的1W频率对齐需指定week_start=6（周日为周结束）
        pl.col("date").dt.truncate("1w", week_start=6).alias("date"),
        "part_of_day"
    )
    .agg(pl.col("user_id").nunique().alias("user_id"))
    .sort("date")
)

print(result)

输出结果与期望一致：

┌─────────────────────┬─────────────┬─────────┐
│ date                ┆ part_of_day ┆ user_id │
│ ---                 ┆ ---         ┆ ---     │
│ datetime[ns]        ┆ cat         ┆ u32     │
╞═════════════════════╪═════════════╪═════════╡
│ 2021-08-01 00:00:00 ┆ night       ┆ 1       │
│ 2021-08-08 00:00:00 ┆ day         ┆ 1       │
│ 2021-08-15 00:00:00 ┆ day         ┆ 1       │
│ 2021-08-29 00:00:00 ┆ day         ┆ 1       │
│ 2021-10-31 00:00:00 ┆ evening     ┆ 1       │
└─────────────────────┴─────────────┴─────────┘

说明 #

• dt.truncate("1w")默认以周一为周起始，若要和Pandas的freq="1W"（周日为周结束）完全对齐，需添加week_start=6参数。
• 分组后通过agg调用nunique()统计唯一用户数，最后排序保证结果顺序符合预期。

方法2：结合group_by与group_by_dynamic（可选） #

如果一定要使用group_by_dynamic，可以先按part_of_day拆分数据，再对每个子数据集做动态分组，最后合并结果：

result = (
    df
    .group_by("part_of_day")
    .map_groups(
        lambda group: group.group_by_dynamic("date", every="1w", week_start=6).agg(pl.col("user_id").nunique())
    )
    .sort("date")
)

但这种方法多了一次分组拆分的开销，不如方法1直接高效，仅作为备选方案。

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