使用Polars实现基于日期范围的DataFrame值求和

需求说明 #

我有两个Polars DataFrame：

• df包含ID、Initial Date、Final Date和Value列
• dates包含需要计算的所有日期
  需求是：对dates中的每个日期，累加df中该日期落在Initial Date到Final Date区间内的所有行的Value值。例如2022-01-01的求和值为10，2022-01-02为10+20，以此类推。

初始代码 #

import polars as pl
from datetime import datetime

data = {
    "ID" : [1, 2, 3, 4, 5],
    "Initial Date" : ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04", "2022-01-05"],
    "Final Date" : ["2022-01-03", "2022-01-06", "2022-01-07", "2022-01-09", "2022-01-07"],
    "Value" : [10, 20, 30, 40, 50]
}

df = pl.DataFrame(data)

dates = pl.datetime_range(
    start=datetime(2022,1,1),
    end=datetime(2022,1,7),
    interval="1d",
    eager = True,
    closed = "both"
    ).to_frame("date")

数据结构展示 #

df的结构 #

shape: (5, 4)
┌─────┬──────────────┬────────────┬───────┐
│ ID  ┆ Initial Date ┆ Final Date ┆ Value │
│ --- ┆ ---          ┆ ---        ┆ ---   │
│ i64 ┆ str          ┆ str        ┆ i64   │
╞═════╪══════════════╪════════════╪═══════╡
│ 1   ┆ 2022-01-01   ┆ 2022-01-03 ┆ 10    │
│ 2   ┆ 2022-01-02   ┆ 2022-01-06 ┆ 20    │
│ 3   ┆ 2022-01-03   ┆ 2022-01-07 ┆ 30    │
│ 4   ┆ 2022-01-04   ┆ 2022-01-09 ┆ 40    │
│ 5   ┆ 2022-01-05   ┆ 2022-01-07 ┆ 50    │
└─────┴──────────────┴────────────┴───────┘

dates的结构 #

shape: (7, 1)
┌─────────────────────┐
│ date                │
│ ---                 │
│ datetime[μs]        │
╞═════════════════════╡
│ 2022-01-01 00:00:00 │
│ 2022-01-02 00:00:00 │
│ 2022-01-03 00:00:00 │
│ 2022-01-04 00:00:00 │
│ 2022-01-05 00:00:00 │
│ 2022-01-06 00:00:00 │
│ 2022-01-07 00:00:00 │
└─────────────────────┴──────────────┘

用户实现代码 #

(
    dates.with_columns(
        pl.sum(
            pl.when(
                (df["Initial Date"] <= pl.col("date")) & (df["Final Date"] >= pl.col("date"))
            ).then(df["Value"]).otherwise(0)
        ).alias("Summed Value")
    )
)

代码正确性分析与优化 #

原代码存在的问题 #

1. 日期类型不匹配：df中的Initial Date和Final Date是字符串类型，而dates的date是datetime类型，直接进行比较会触发类型错误，无法得到正确结果。
2. 性能瓶颈：该实现属于暴力比对逻辑，每个日期都要遍历df的所有行进行条件判断，当数据量较大时（比如df有上万行、dates有上千个日期），时间复杂度会达到O(N*M)，性能会急剧下降。

优化方案：事件法（差分数组思想） #

针对区间求和类问题，最高效的解法是利用差分数组思想：将每个区间的起始日期标记为「加上对应Value」，结束日期的次日标记为「减去对应Value」，最后对日期排序后做累加，即可快速得到每个日期的总和。这种方法的时间复杂度为O(N+M)，性能远超暴力比对。

优化后的代码

# 第一步：统一日期类型，将df中的字符串日期转为datetime
df = df.with_columns(
    pl.col("Initial Date", "Final Date").str.to_datetime()
)

# 第二步：生成事件数据：起始日期加Value，结束日期次日减Value
events = pl.concat([
    # 区间开始事件：加Value
    df.select(
        pl.col("Initial Date").alias("date"),
        pl.col("Value").alias("delta")
    ),
    # 区间结束事件：减Value（结束日期的次日）
    df.select(
        (pl.col("Final Date") + pl.duration(days=1)).alias("date"),
        (-pl.col("Value")).alias("delta")
    )
])

# 第三步：合并日期与事件，排序后累加得到结果
result = (
    dates.join(events, on="date", how="outer")
    .sort("date")
    .fill_null(0)  # 没有事件的日期delta为0
    .with_columns(pl.col("delta").cumsum().alias("Summed Value"))
    .filter(pl.col("date").is_in(dates["date"]))  # 只保留原dates中的日期
    .select("date", "Summed Value")
)

print(result)

代码说明

• 类型统一：先将df的日期列转为datetime，避免类型不匹配问题。
• 事件生成：每个区间对应两个事件，确保区间内的所有日期都会包含该Value的贡献。
• 累加计算：合并所有日期和事件后排序，通过累加delta得到每个日期的总Value，最后过滤出原dates中的日期即可得到目标结果。

验证结果

运行后得到的结果符合预期：

shape: (7, 2)
┌─────────────────────┬──────────────┐
│ date                ┆ Summed Value │
│ ---                 ┆ ---          │
│ datetime[μs]        ┆ i64          │
╞═════════════════════╪══════════════╡
│ 2022-01-01 00:00:00 ┆ 10           │
│ 2022-01-02 00:00:00 ┆ 30           │
│ 2022-01-03 00:00:00 ┆ 60           │
│ 2022-01-04 00:00:00 ┆ 100          │
│ 2022-01-05 00:00:00 ┆ 140          │
│ 2022-01-06 00:00:00 ┆ 90           │
│ 2022-01-07 00:00:00 ┆ 80           │
└─────────────────────┴──────────────┘

原代码的修正版本（仅适合小数据量） #

如果一定要沿用原思路，需要先修正日期类型问题，代码如下：

# 先统一日期类型
df = df.with_columns(
    pl.col("Initial Date", "Final Date").str.to_datetime()
)

# 修正后的原逻辑代码
result = dates.with_columns(
    pl.sum(
        pl.when(
            (pl.col("date") >= df["Initial Date"]) & (pl.col("date") <= df["Final Date"])
        ).then(df["Value"]).otherwise(0)
    ).alias("Summed Value")
)

print(result)

但此方法仅适合数据量较小的场景，数据量大时推荐使用事件法优化方案。

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