如何用Polars按多段折扣时间区间过滤DataFrame？

问题背景 #

现有两个Polars DataFrame：

• discount_dates：存储每个产品的多段独立折扣起止日期
• data_product：存储产品的每日销量数据

需求是从data_product中筛选出日期落在对应产品任意一段折扣区间内的行。原实现通过聚合产品折扣的最小开始日期和最大结束日期来过滤，忽略了折扣区间之间的间隙，导致错误包含非折扣日期的行。

示例数据 #

折扣日期表 discount_dates #

import polars as pl

discount_dates = pl.DataFrame({
    "product_id": ["A", "A", "B"],
    "discount_start": ["2023-01-01", "2023-01-10", "2023-01-05"],
    "discount_end": ["2023-01-05", "2023-01-15", "2023-01-20"]
}).with_columns([
    pl.col("discount_start").str.to_date(),
    pl.col("discount_end").str.to_date()
])

销量数据表 data_product #

data_product = pl.DataFrame({
    "product_id": ["A", "A", "A", "B", "B"],
    "sale_date": ["2023-01-03", "2023-01-07", "2023-01-12", "2023-01-03", "2023-01-10"],
    "sales": [100, 200, 150, 300, 250]
}).with_columns(pl.col("sale_date").str.to_date())

期望结果 #

仅保留落在有效折扣区间内的行，即：

原错误代码及问题分析 #

# 错误实现：聚合最小/最大日期，忽略区间间隙
wrong_result = data_product.join(
    discount_dates.group_by("product_id").agg(
        pl.col("discount_start").min().alias("min_start"),
        pl.col("discount_end").max().alias("max_end")
    ),
    on="product_id"
).filter(
    pl.col("sale_date").is_between(pl.col("min_start"), pl.col("max_end"))
)

问题：该方法将产品的所有折扣区间合并为一个从最早开始到最晚结束的大区间，忽略了区间之间的间隙（比如产品A的折扣区间是2023-01-01~05和2023-01-10~15，中间2023-01-06~09无折扣），导致错误包含间隙内的日期（如示例中A的2023-01-07）。

正确实现方式 #

方法1：使用exists子查询（推荐，简洁高效） #

直接在filter中判断当前行的日期是否存在匹配的产品折扣区间，无需额外去重：

correct_result = data_product.filter(
    pl.exists(
        discount_dates,
        (pl.col("product_id") == discount_dates["product_id"]) &
        pl.col("sale_date").is_between(discount_dates["discount_start"], discount_dates["discount_end"])
    )
)

方法2：左连接后过滤（适合需要保留折扣区间信息的场景） #

如果需要同时保留折扣区间的详情，可以先做左连接，再筛选出日期落在区间内的行，最后去重（避免同一行匹配多个折扣区间导致重复）：

correct_result = data_product.join(
    discount_dates,
    on="product_id",
    how="left"
).filter(
    pl.col("sale_date").is_between(pl.col("discount_start"), pl.col("discount_end"))
).select(data_product.columns).unique()

原理说明 #

两种方法都精准匹配每一段折扣区间：

• 方法1通过exists子查询，对data_product的每一行检查是否存在对应的折扣区间覆盖该日期，逻辑简洁且性能优异。
• 方法2通过左连接关联所有可能的折扣区间，再筛选出日期落在区间内的行，适合需要进一步处理折扣区间信息的场景。

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