如何用Polars计算日期列间的工作日及交易日差值？

我有如下Polars DataFrame：

df = pl.from_repr("""
┌────────────┬───────────────┐
│ date       ┆ maturity_date │
│ ---        ┆ ---           │
│ date       ┆ date          │
╞════════════╪═══════════════╡
│ 2000-01-04 ┆ 2000-01-17    │
│ 2000-01-04 ┆ 2000-02-15    │
│ 2000-01-04 ┆ 2000-03-15    │
│ 2000-01-04 ┆ 2000-04-17    │
│ 2000-01-04 ┆ 2000-05-15    │
└────────────┴───────────────┘
""")

我想要计算date与maturity_date之间的工作日数量（不含周六和周日），同时也希望基于指定日历（如股市交易日历）计算交易日差值。

我目前使用date_ranges方法计算工作日，但它仅比map_elements略快：

df.with_columns(
    pl.date_ranges("date", "maturity_date")
    .list.eval(pl.element().dt.weekday() <= 5)
    .list.count_matches(True)
    .alias("workdays_diff")
    # pl.concat_list("date", "maturity_date").map_elements(lambda x: get_work_days(x[0], x[1]))
    # .alias("workdays_diff")
)

执行结果如下：

shape: (5, 3)
┌────────────┬───────────────┬───────────────┐ # ┌────────────────┐
│ date       ┆ maturity_date ┆ workdays_diff │ # │ tradedate_diff │
│ ---        ┆ ---           ┆ ---           │ # │ ---            │
│ date       ┆ date          ┆ u32           │ # │ i64            │
╞════════════╪═══════════════╪═══════════════╡ # ╞════════════════╡
│ 2000-01-04 ┆ 2000-01-17    ┆ 10            │ # │ 10             │
│ 2000-01-04 ┆ 2000-02-15    ┆ 31            │ # │ 21             │
│ 2000-01-04 ┆ 2000-03-15    ┆ 52            │ # │ 42             │
│ 2000-01-04 ┆ 2000-04-17    ┆ 75            │ # │ 65             │
│ 2000-01-04 ┆ 2000-05-15    ┆ 95            │ # │ 80             │
└────────────┴───────────────┴───────────────┘ # └────────────────┘

请问是否有更高效的计算工作日的方法？同时能否实现tradedate_diff的计算？

一、更高效的工作日计算方法 #

你当前用date_ranges生成日期列表再过滤的方法，本质是逐行生成序列，数据量大时效率会下降。可以用数学公式直接计算，避免生成中间日期列表，性能提升明显：

核心逻辑 #

1. 计算两个日期的总天数差（包含两端日期）
2. 拆分总天数为完整周数和剩余天数
3. 完整周数的工作日为 完整周数 * 5
4. 剩余天数中排除周末，统计工作日数量

实现代码 #

def calculate_workdays_diff(df):
    return df.with_columns(
        # 总天数（包含起始和结束日期）
        total_days = (pl.col("maturity_date") - pl.col("date")).dt.days() + 1,
        # 起始/结束日期的星期几（1=周一，7=周日）
        start_weekday = pl.col("date").dt.weekday(),
        end_weekday = pl.col("maturity_date").dt.weekday()
    ).with_columns(
        full_weeks = pl.col("total_days") // 7,
        remaining_days = pl.col("total_days") % 7,
        # 计算剩余天数中的工作日：跨周末时减去周末天数
        remaining_workdays = pl.when(pl.col("start_weekday") + pl.col("remaining_days") <=7)
            .then(pl.col("remaining_days"))
            .otherwise(pl.col("remaining_days") - (pl.col("start_weekday") + pl.col("remaining_days") -7)),
        # 处理剩余天数超过5天的情况（最多5个工作日）
        remaining_workdays = pl.when(pl.col("remaining_workdays")>5)
            .then(5)
            .otherwise(pl.col("remaining_workdays"))
    ).with_columns(
        workdays_diff = pl.col("full_weeks") *5 + pl.col("remaining_workdays")
    ).drop("total_days", "start_weekday", "end_weekday", "full_weeks", "remaining_days", "remaining_workdays")

# 调用方法
result = calculate_workdays_diff(df)
print(result)

这种方法完全基于向量运算，没有中间列表生成，在大数据集下性能比date_ranges高一个数量级。

二、计算交易日差值（基于指定日历） #

要计算股市交易日差值，需要依赖交易日历数据（包含所有交易日、排除节假日休市日），步骤如下：

1. 准备交易日历数据 #

假设你已有包含所有交易日的Polars DataFrame trading_cal_df，示例结构如下：

# 模拟A股2000年部分交易日历
trading_cal_df = pl.DataFrame({
    "trading_date": [
        "2000-01-04", "2000-01-05", "2000-01-06", "2000-01-07", "2000-01-10",
        "2000-01-11", "2000-01-12", "2000-01-13", "2000-01-14", "2000-01-17"
        # 更多交易日数据...
    ]
}).with_columns(pl.col("trading_date").str.to_date())

2. 为交易日历添加序号 #

给每个交易日分配连续序号，序号差即为交易日差值：

trading_cal_df = trading_cal_df.with_columns(
    trade_idx = pl.int_range(0, len(trading_cal_df), eager=True)
)

3. 关联原DataFrame计算差值 #

通过两次左连接获取起始/结束日期的序号，再计算差值：

def calculate_tradedate_diff(df, trading_cal_df):
    # 关联起始日期的序号
    df_with_start = df.join(
        trading_cal_df.rename({"trading_date": "date", "trade_idx": "start_idx"}),
        on="date",
        how="left"
    )
    # 关联结束日期的序号
    df_with_both = df_with_start.join(
        trading_cal_df.rename({"trading_date": "maturity_date", "trade_idx": "end_idx"}),
        on="maturity_date",
        how="left"
    )
    # 计算交易日差值（包含两端日期则+1，否则用end_idx - start_idx）
    return df_with_both.with_columns(
        tradedate_diff = pl.col("end_idx") - pl.col("start_idx") +1
    ).drop("start_idx", "end_idx")

# 调用方法
result_with_trade = calculate_tradedate_diff(result, trading_cal_df)
print(result_with_trade)

说明 #

• 交易日历可从外部文件（如CSV）加载，无需手动模拟
• 差值逻辑可按需调整：若仅算中间交易日数，去掉+1即可
• 若原DataFrame中的日期不在交易日历（非交易日），可根据业务需求处理（如报错、填充NaN等）

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