如何基于左pl.LazyFrame列内容实现两表两列条件连接？

需求说明 #

需要对两个Polars LazyFrame（lf1和lf2）执行条件左连接：

• 当lf1的col_2不为空时，用col_2关联lf2的col_a
• 当lf1的col_3不为空（且col_2为空）时，用col_3关联lf2的col_b
• 若两列都为空，则不进行关联，保留原行且col_c为null

原始数据代码 #

import polars as pl

lf1 = pl.LazyFrame(
    data={
        "col_1": ["a", "b", "c"],
        "col_2": ["d", None, None],
        "col_3": [None, "e", None],
    },
)

lf2 = pl.LazyFrame(
    data={
        "col_a": ["d", "xyz"],
        "col_b": ["xyz", "e"],
        "col_c": ["relevant_info_1", "relevant_info_2"],
    },
)  

期望逻辑伪代码 #

lf1.join(lf2,
  when(col("col_2").is_not_null().then(left_on="col_2", right_on="col_a")
  when(col("col_3").is_not_null().then(left_on="col_3", right_on="col_b")
  otherwise(do_nothing)
)

实现方案 #

Polars原生join不支持动态指定连接键，我们可以按条件拆分数据为三个互斥子集，分别处理后合并：

# 情况1：col_2不为空，按col_2-col_a连接
case1 = lf1.filter(pl.col("col_2").is_not_null()).join(lf2, left_on="col_2", right_on="col_a", how="left")

# 情况2：col_2为空但col_3不为空，按col_3-col_b连接
case2 = lf1.filter(pl.col("col_2").is_null() & pl.col("col_3").is_not_null()).join(lf2, left_on="col_3", right_on="col_b", how="left")

# 情况3：col_2和col_3都为空，直接添加null的col_c列
case3 = lf1.filter(pl.col("col_2").is_null() & pl.col("col_3").is_null()).with_columns(col_c=None)

# 合并结果并按col_1排序
result = pl.concat([case1, case2, case3]).sort("col_1").collect()
print(result)

预期输出 #

shape: (3, 4)
┌───────┬───────┬───────┬─────────────────┐
│ col_1 ┆ col_2 ┆ col_3 ┆ col_c           │
│ ---   ┆ ---   ┆ ---   ┆ ---             │
│ str   ┆ str   ┆ str   ┆ str             │
╞═══════╪═══════╪═══════╪═════════════════╡
│ a     ┆ d     ┆ null  ┆ relevant_info_1 │
│ b     ┆ null  ┆ e     ┆ relevant_info_2 │
│ c     ┆ null  ┆ null  ┆ null            │
└───────┴───────┴───────┴─────────────────┘

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