如何在Polars中基于ID及日期区间条件实现左连接？

问题描述 #

我有两个Polars DataFrame：

• df_1：包含id、start、end、value列，其中start和end定义日期区间
• df_2：包含id、dates列

需要以df_2为左表执行左连接，连接条件为：

1. 两张表的id完全匹配
2. df_2的dates值落在df_1对应id的start与end区间内（包含边界）

示例数据 #

import polars as pl
from datetime import date

df_1 = pl.DataFrame(
    {
        "id": ["abc", "abc", "456"],
        "start": [date(2022, 1, 1), date(2022, 3, 4), date(2022, 5, 11)],
        "end": [date(2022, 2, 4), date(2022, 3, 10), date(2022, 5, 16)],
        "value": [10, 3, 4]
    }
)

df_2 = pl.DataFrame(
    {
        "id": ["abc", "abc", "456", "abc", "abc", "456"],
        "dates": [date(2022, 1, 2), date(2022, 3, 4), date(2022, 5, 11), date(2022, 1, 4), date(2022, 3, 7), date(2022, 5, 13)],
    }
)

预期结果 #

shape: (6, 3)
┌─────┬────────────┬───────┐
│ id  ┆ dates      ┆ value │
│ --- ┆ ---        ┆ ---   │
│ str ┆ date       ┆ i64   │
╞═════╪════════════╪═══════╡
│ abc ┆ 2022-01-02 ┆ 10    │
│ abc ┆ 2022-03-04 ┆ 3     │
│ 456 ┆ 2022-05-11 ┆ 4     │
│ abc ┆ 2022-01-04 ┆ 10    │
│ abc ┆ 2022-03-07 ┆ 3     │
│ 456 ┆ 2022-05-13 ┆ 4     │
└─────┴────────────┴───────┘

解决方案 #

Polars支持在join方法中同时指定等值匹配列和非等值条件，直接实现需求：

result = df_2.join(
    df_1,
    on="id",  # 等值匹配id列
    condition=(pl.col("dates") >= pl.col("start")) & (pl.col("dates") <= pl.col("end")),  # 日期区间判断
    how="left"  # 以df_2为左表的左连接
).select(["id", "dates", "value"])  # 保留需要的列

print(result)

代码说明 #

1. on="id"：指定两张表通过id列进行等值匹配
2. condition参数：添加非等值的日期区间判断，确保dates落在start和end之间（包含边界值）
3. how="left"：确保左表df_2的所有行都被保留，匹配到的行补充value列，未匹配的行value会显示为null（本示例中所有行都能匹配到）
4. select方法：筛选出最终需要的三列，与预期结果一致

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