如何在 Polars 中动态配置周末掩码实现业务日加法-Python教程-PHP中文网

如何在 Polars 中动态配置周末掩码实现业务日加法

本文介绍在 Polars 中为每条记录动态指定周末（WKND1/WKND2）并精确计算业务日后新日期的方法，通过 partition_by 分组 + add_business_days 配合自定义 week_mask 实现灵活、高效、可扩展的日期运算。

本文介绍在 polars 中为每条记录动态指定周末（wknd1/wknd2）并精确计算业务日后新日期的方法，通过 `partition_by` 分组 + `add_business_days` 配合自定义 `week_mask` 实现灵活、高效、可扩展的日期运算。

在 Polars 中，add_business_days() 方法默认仅支持全局固定的 week_mask（如 (True, True, True, True, True, False, False) 表示仅周六日休息），但实际业务中，不同地区、组织或合同可能定义不同的“非工作日”——例如某客户将周五和周六设为周末，另一客户仅将周六设为休息日。此时，无法直接传入表达式（如 pl.col('W_MASK')）作为 week_mask 参数，因为该参数要求是 Python 原生 Sequence[bool]（如 tuple/list），而非 Polars Expr。

解决方案是采用 分组—逐组构造掩码—分别计算—合并结果 的策略，核心步骤如下：

✅ 步骤一：定义标准星期顺序

weekdays = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']

该顺序严格对应 week_mask 的 7 元素布尔序列：索引 0 → 周一是否为工作日（True），索引 6 → 周日是否为工作日。

✅ 步骤二：按周末配置分组处理

使用 df.partition_by('WKND1', 'WKND2', as_dict=True) 将数据按 WKND1 和 WKND2 的唯一组合切分为多个子 DataFrame。每个键 k 是一个元组（如 ('Friday', 'Saturday')），值 v 是对应子集。

✅ 步骤三：为每组动态生成 `week_mask`

对每组键 k，构造布尔列表：

week_mask = [day not in k for day in weekdays]

注意：k 是元组（可能含 None），in 操作安全；None 会被自动忽略（因 None not in ('Saturday', None) 为 True），故单周末列（如 WKND2=null）也能正确处理。

✅ 步骤四：应用业务日加法并合并

pl.concat([
    v.with_columns(
        pl.col('DT').dt.add_business_days(
            pl.col('N_DAYS'),
            week_mask=[day not in k for day in weekdays]
        ).alias('NEW_DT')
    )
    for k, v in df.partition_by('WKND1', 'WKND2', as_dict=True).items()
]).sort('DT')

⚠️ 注意事项

partition_by(..., as_dict=True) 在 Polars ≥ 0.20.0 可用；旧版本请改用 group_by(...).agg(...) + explode() 组合模拟。
若存在大量唯一周末组合（如数百种），分组数过多可能影响性能；此时建议预聚合或改用 UDF + map_rows（但会失去惰性执行优势）。
week_mask 中 True 表示该 weekday 是工作日，False 表示休息日——逻辑易反，务必验证首例：('Saturday','Sunday') → 周六/日非工作日 → 掩码末两位为 False → [T,T,T,T,T,F,F] ✔️