如何在Polars中使用多条件（等式+比较）连接两个DataFrame？

问题背景 #

需将包含三类条件的SQL左连接转换为Polars代码，三个连接条件为：

• tab1.article_no = tab2.article_no（等值匹配）
• DIV(tab1.variant_no,1000) = tab2.variant_no（计算后等值匹配）
• tab1.date_of_day BETWEEN tab2.date_from AND tab2.date_to（日期范围匹配）

直接通过前两个等值条件连接后再过滤日期，会导致百万级数据场景下内存占用过高，因此需要在左连接阶段直接应用全部条件。

解决方案 #

利用Polars的join方法结合动态计算列和**condition参数**，在连接过程中同时校验所有条件，避免生成不必要的中间行：

1. 准备示例数据 #

import polars as pl

# 表1数据
table_1 = pl.DataFrame({
    "article_no": [1, 1, 2, 3],
    "variant_no": [1234, 5678, 9012, 4321],
    "date_of_day": ["2023-01-15", "2023-02-20", "2023-03-10", "2023-04-01"]
}).with_columns(pl.col("date_of_day").str.to_date())

# 表2数据
table_2 = pl.DataFrame({
    "article_no": [1, 1, 2],
    "variant_no": [1, 5, 9],
    "date_from": ["2023-01-01", "2023-02-01", "2023-03-01"],
    "date_to": ["2023-01-31", "2023-02-28", "2023-03-31"],
    "extra_data": ["A", "B", "C"]
}).with_columns(pl.col("date_from", "date_to").str.to_date())

2. 多条件左连接实现 #

result = table_1.join(
    table_2,
    # 处理前两个等值条件：直接在left_on中计算DIV逻辑，无需提前生成派生列
    left_on=["article_no", pl.col("variant_no") // 1000],
    right_on=["article_no", "variant_no"],
    how="left",
    # 添加日期范围的非等值条件，在连接阶段直接过滤
    condition=(pl.col("date_of_day") >= pl.col("date_from")) & (pl.col("date_of_day") <= pl.col("date_to"))
)

关键说明 #

• 直接在left_on中计算variant_no // 1000，等价于SQL中的DIV(tab1.variant_no,1000)，无需额外预处理列
• condition参数会在连接过程中直接过滤不满足日期范围的行，避免先生成大量冗余连接结果再过滤，大幅降低内存开销
• 左连接逻辑保留table_1所有行，不满足条件的table_2字段会填充null

最终输出 #

shape: (4, 7)
┌────────────┬────────────┬─────────────┬────────────┬────────────┬────────────┬────────────┐
│ article_no ┆ variant_no ┆ date_of_day ┆ variant_no ┆ date_from  ┆ date_to    ┆ extra_data │
│ ---        ┆ ---        ┆ ---         ┆ ---        ┆ ---        ┆ ---        ┆ ---        │
│ i64        ┆ i64        ┆ date        ┆ i64        ┆ date       ┆ date       ┆ str        │
╞════════════╪════════════╪═════════════╪════════════╪════════════╪════════════╪════════════╡
│ 1          ┆ 1234       ┆ 2023-01-15  ┆ 1          ┆ 2023-01-01 ┆ 2023-01-31 ┆ A          │
│ 1          ┆ 5678       ┆ 2023-02-20  ┆ 5          ┆ 2023-02-01 ┆ 2023-02-28 ┆ B          │
│ 2          ┆ 9012       ┆ 2023-03-10  ┆ 9          ┆ 2023-03-01 ┆ 2023-03-31 ┆ C          │
│ 3          ┆ 4321       ┆ 2023-04-01  ┆ null       ┆ null       ┆ null       ┆ null       │
└────────────┴────────────┴─────────────┴────────────┴────────────┴────────────┴────────────┘

内容的提问来源于stack exchange，提问作者Solomon Papathoti Leo