如何用Polars表达式按日按类别计算累计行计数及扩展需求

原始数据 #

给定如下Polars DataFrame：

import datetime
import polars as pl

df = pl.DataFrame(
    {
        "time": [
            datetime.datetime(2023, 1, 1, 9),
            datetime.datetime(2023, 1, 1, 10),
            datetime.datetime(2023, 1, 1, 12),
            datetime.datetime(2023, 1, 2, 9),
            datetime.datetime(2023, 1, 2, 10),
            datetime.datetime(2023, 1, 3, 12),
        ],
        "category": [1,1,2,1,2,1],
    }
)

需求目标 #

构造表达式row_count_by_day，执行以下代码后，生成的列需表示每个类别在单日的累计行计数：

expr = ...alias("row_count_by_day")
df = df.with_columns(expr)
print(df)

预期输出：

shape: (6, 3)
┌─────────────────────┬──────────┬──────────────────┐
│ time                ┆ category ┆ row_count_by_day │
│ ---                 ┆ ---      ┆ ---              │
│ datetime[μs]        ┆ i64      ┆ i64              │
╞═════════════════════╪══════════╪══════════════════╡
│ 2023-01-01 09:00:00 ┆ 1        ┆ 1                │
│ 2023-01-01 10:00:00 ┆ 1        ┆ 2                │
│ 2023-01-01 12:00:00 ┆ 2        ┆ 1                │
│ 2023-01-02 09:00:00 ┆ 1        ┆ 1                │
│ 2023-01-02 10:00:00 ┆ 2        ┆ 1                │
│ 2023-01-03 12:00:00 ┆ 1        ┆ 1                │
└─────────────────────┴──────────┴──────────────────┘

同时需要扩展至多类别场景：例如当某类别在另一列存在重复值时，相同值仅计数一次。

基础场景解决方案 #

利用Polars的窗口函数row_number()，按日期（从time字段提取）+category分区，按time排序后生成累计序号：

expr = pl.row_number().over(
    [pl.col("time").dt.date(), "category"]
).sort_by("time").alias("row_count_by_day")

df = df.with_columns(expr)
print(df)

原理说明 #

• pl.col("time").dt.date()：从datetime类型的time字段提取日期部分，作为分组的时间维度
• over([日期, category])：指定窗口的分区键，确保计数仅在同一日期、同一类别内进行
• row_number()：在每个分区内按time顺序生成递增序号，即单日单类别的累计行计数

扩展场景：去重后累计计数 #

如果需要对额外字段（如sub_category）去重，同一日期+类别下相同值仅计数一次，可通过以下方式实现：

示例扩展数据 #

df_extended = pl.DataFrame(
    {
        "time": [
            datetime.datetime(2023, 1, 1, 9),
            datetime.datetime(2023, 1, 1, 10),
            datetime.datetime(2023, 1, 1, 12),
            datetime.datetime(2023, 1, 1, 13),
            datetime.datetime(2023, 1, 2, 9),
        ],
        "category": [1,1,1,2,1],
        "sub_category": [10,10,20,30,10]
    }
)

构造表达式 #

expr_extended = (
    pl.when(
        pl.col("sub_category").is_first().over([pl.col("time").dt.date(), "category", "sub_category"])
    )
    .then(1)
    .otherwise(0)
    .cum_sum()
    .over([pl.col("time").dt.date(), "category"])
    .sort_by("time")
    .alias("unique_row_count_by_day")
)

df_extended = df_extended.with_columns(expr_extended)
print(df_extended)

输出结果 #

shape: (5, 4)
┌─────────────────────┬──────────┬───────────────┬───────────────────────┐
│ time                ┆ category ┆ sub_category  ┆ unique_row_count_by_day │
│ ---                 ┆ ---      ┆ ---           ┆ ---                   │
│ datetime[μs]        ┆ i64      ┆ i64           ┆ i64                   │
╞═════════════════════╪══════════╪═══════════════╪═══════════════════════╡
│ 2023-01-01 09:00:00 ┆ 1        ┆ 10            ┆ 1                     │
│ 2023-01-01 10:00:00 ┆ 1        ┆ 10            ┆ 1                     │
│ 2023-01-01 12:00:00 ┆ 1        ┆ 20            ┆ 2                     │
│ 2023-01-01 13:00:00 ┆ 2        ┆ 30            ┆ 1                     │
│ 2023-01-02 09:00:00 ┆ 1        ┆ 10            ┆ 1                     │
└─────────────────────┴──────────┴───────────────┴───────────────────────┘

原理说明 #

1. pl.col("sub_category").is_first().over([日期, category, sub_category])：判断当前行的sub_category在“日期+类别+子类别”分组中是否是首次出现，首次出现标记为1，否则为0
2. cum_sum().over([日期, category])：在“日期+类别”窗口内对标记值做累计求和，得到去重后的累计计数

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