基于Polars的子查询与阈值首次触发的复杂操作需求

数据定义 #

df_original代表A、B站点的冰淇淋检查员数据：

df_original = pl.DataFrame(
    {
        "station": ["A", "A", "A", "A", "B", "B", "B", "B"],
        "ice_cream_date": [1, 2, 3, 4, 1, 2, 3, 4],
        "customers": [10, 20, 30, 5, 5, 7, 4, 10],
        "event": [0, 1, 0, 1, 1, 0, 1, 0],
    }
)

• ice_cream_date：编码日期
• customers：客户数量
• event：检查员检查事件的二进制编码

df_events是df_original中event == 1的数据集：

df_events =  pl.DataFrame(
    {
        "events": [1, 1, 1, 1],
        "ice_cream_date": [2, 4, 1, 3],
        "station": ["A", "A", "B", "B"],
        "customers": [20, 5, 5, 4],
        "evaluation_span": [10, 2, 2, 2],
        "evaluation_end_date": [4, 4, 3, 4],
        "evaluation_end_date_customer": [5, 5, 4, 10],
        "good_customer": [30, 10, 8, 12],
        "bad_customer": [5, 0, 4, 3],
    }
)

• evaluation_span：检查日期后的天数
• evaluation_end_date：检查最晚日期（若日期+evaluation_span超过最大可用日期，则取最大可用日期）
• evaluation_end_date_customer：df_original中evaluation_end_date对应的客户数
• good_customer：检查员给站点"良好"评级的客户数阈值
• bad_customer：检查员给站点"不良"评级的客户数阈值

需求说明 #

需为每个非零事件标记站点获得的是good还是bad评级，规则如下：
若检查员发现客户数首次超过良好阈值（或低于不良阈值），则无论评估周期长度，直接给予对应评级。

预期输出 #

shape: (4, 7)
┌────────────────┬─────────┬────────────────────┬──────┬─────┬────────────────┬────────────────────┐
│ ice_cream_date ┆ station ┆ evaluation_custome ┆ good ┆ bad ┆ evaluation_end ┆ actual_evaluation_ │
│ ---            ┆ ---     ┆ rs                 ┆ ---  ┆ --- ┆ ---            ┆ end_date           │
│ i64            ┆ str     ┆ ---                ┆ i64  ┆ i64 ┆ i64            ┆ ---                │
│                ┆         ┆ i64                ┆      ┆     ┆                ┆ i64                │
╞════════════════╪═════════╪════════════════════╪══════╪═════╪════════════════╪════════════════════╡
│ 2              ┆ A       ┆ 30                 ┆ 1    ┆ 0   ┆ 0              ┆ 3                  │
│ 4              ┆ A       ┆ 5                  ┆ 0    ┆ 0   ┆ 1              ┆ 4                  │
│ 1              ┆ B       ┆ 4                  ┆ 0    ┆ 1   ┆ 1              ┆ 3                  │
│ 3              ┆ B       ┆ 10                 ┆ 0    ┆ 0   ┆ 1              ┆ 4                  │
└────────────────┴─────────┴────────────────────┴──────┴─────┴────────────────┴────────────────────┘

• evaluation_customers：评估客户数
• good：二进制标记，表示在(事件日期, 事件日期+评估周期]区间内该检查事件是否给予良好评级
• bad：二进制标记，表示在上述区间内该检查事件是否给予不良评级
• evaluation_end：未达标标记，表示在评估区间内客户数未达到良好/不良阈值
• actual_evaluation_end_date：实际评估结束日期（达到阈值的日期）

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