如何在Polars中实现带锚定行条件的分组滚动计数

核心需求 #

按指定分组统计每一行对应的、满足以下条件的记录数：

• days字段值大于90
• 与当前行同属一个分组
• Date Closed日期早于当前行的Due Date

实现方案 #

1. 准备数据（确保日期类型正确）

首先需要将日期字符串转为Polars的date类型，否则无法进行日期比较：

import polars as pl

# 示例数据
df = pl.DataFrame({
    "group": ["A", "A", "A", "B", "B"],
    "days": [100, 80, 120, 95, 85],
    "Date Closed": ["2023-01-01", "2023-02-01", "2023-03-01", "2023-01-15", "2023-02-15"],
    "Due Date": ["2023-02-15", "2023-03-01", "2023-04-01", "2023-02-20", "2023-03-20"]
}).with_columns(
    pl.col("Date Closed").str.to_date(),
    pl.col("Due Date").str.to_date()
)

2. 分组统计逻辑

利用group_by结合map_groups，在每个分组内通过广播对比实现逐行统计：

df = df.group_by("group").map_groups(
    lambda group_df: group_df.with_columns(
        # 广播对比：每个Due Date与组内所有Date Closed比较，结合days>90条件后按行求和
        ((group_df["days"] > 90) & (group_df["Date Closed"] < group_df["Due Date"].to_numpy()[:, None])).sum(axis=1).alias("overdue_90_plus_count")
    )
)

代码解释 #

• group_by("group"): 按分组字段拆分数据
• map_groups: 对每个分组的子DataFrame单独处理
• group_df["Due Date"].to_numpy()[:, None]: 将当前分组的Due Date转为列向量，实现与组内所有Date Closed的逐行广播对比
• (group_df["days"] > 90) & (...): 生成布尔矩阵，标记满足双重条件的记录
• .sum(axis=1): 按行求和，得到当前行对应的符合条件的记录数

输出结果 #

执行后会新增overdue_90_plus_count列，示例输出如下：

shape: (5, 5)
┌───────┬──────┬──────────────┬──────────┬────────────────────────┐
│ group ┆ days ┆ Date Closed  ┆ Due Date ┆ overdue_90_plus_count │
│ ---   ┆ ---  ┆ ---          ┆ ---      ┆ ---                    │
│ str   ┆ i64  ┆ date         ┆ date     ┆ u32                    │
╞═══════╪══════╪══════════════╪══════════╪════════════════════════╡
│ A     ┆ 100  ┆ 2023-01-01   ┆ 2023-02-15 ┆ 1                      │
│ A     ┆ 80   ┆ 2023-02-01   ┆ 2023-03-01 ┆ 1                      │
│ A     ┆ 120  ┆ 2023-03-01   ┆ 2023-04-01 ┆ 2                      │
│ B     ┆ 95   ┆ 2023-01-15   ┆ 2023-02-20 ┆ 1                      │
│ B     ┆ 85   ┆ 2023-02-15   ┆ 2023-03-20 ┆ 1                      │
└───────┴──────┴──────────────┴──────────┴────────────────────────┘

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