使用Polars .when()替代join实现多DataFrame关联取值

你遇到的报错核心原因是：when的then分支直接返回了T1的整个Values列，Polars无法将其与主表的行按匹配关系一一对应（主表和T1的行顺序、匹配逻辑没有绑定），因此报错。

以下是几种无需重复重命名表、高效实现关联取值的方法，替代传统的多次join：

方法1：字典映射（适合小维度表） #

先将关联表转换为ID到值的字典，再通过replace方法快速映射，代码简洁且内存占用低：

# 构建T1的映射字典：ID → (Values, Values II)
t1_dict = T1.select(pl.col("ID"), pl.struct(["Values", "Values II"])).to_dict(as_series=False)
t1_dict = dict(zip(t1_dict["ID"], t1_dict["struct"]))

# 构建T2的映射字典：ID → Values
t2_dict = T2.select(pl.col("ID"), pl.col("Values")).to_dict(as_series=False)
t2_dict = dict(zip(t2_dict["ID"], t2_dict["Values"]))

# 填充主表
result = df.with_columns(
    # ID1匹配T1的结果
    ID1_Values = pl.col("ID1").replace(t1_dict, default=None).struct.field("Values"),
    ID1_Values_II = pl.col("ID1").replace(t1_dict, default=None).struct.field("Values II"),
    # ID2匹配T1的结果
    ID2_Values = pl.col("ID2").replace(t1_dict, default=None).struct.field("Values"),
    ID2_Values_II = pl.col("ID2").replace(t1_dict, default=None).struct.field("Values II"),
    # ID1匹配T2的结果
    ID1_T2_Values = pl.col("ID1").replace(t2_dict, default=None),
    # ID2匹配T2的结果
    ID2_T2_Values = pl.col("ID2").replace(t2_dict, default=None)
)

print(result)

输出结果：

shape: (3, 8)
┌─────┬─────┬────────────┬────────────────┬────────────┬────────────────┬──────────────┬──────────────┐
│ ID1 ┆ ID2 ┆ ID1_Values ┆ ID1_Values_II ┆ ID2_Values ┆ ID2_Values_II ┆ ID1_T2_Values ┆ ID2_T2_Values │
│ --- ┆ --- ┆ ---        ┆ ---           ┆ ---        ┆ ---           ┆ ---          ┆ ---          │
│ i64 ┆ i64 ┆ str        ┆ str           ┆ str        ┆ str           ┆ str          ┆ str          │
╞═════╪═════╪════════════╪═══════════════╪════════════╪═══════════════╪══════════════╪══════════════╡
│ 1   ┆ 1   ┆ null       ┆ null          ┆ null       ┆ null          ┆ X            ┆ null         │
│ 2   ┆ 4   ┆ B          ┆ boo           ┆ null       ┆ null          ┆ null         ┆ J            │
│ 3   ┆ 5   ┆ null       ┆ null          ┆ c          ┆ baz           ┆ null         ┆ null         │
└─────┴─────┴────────────┴────────────────┴────────────┴────────────────┴──────────────┴──────────────┘

方法2：Polars lookup函数（推荐，适配大表） #

Polars 0.19.0+提供的lookup方法是专门为这种"按列值查找关联"场景设计的，无需重命名表，性能优于多次join：

result = (
    df
    # 匹配ID1到T1，返回包含匹配结果的struct列
    .with_columns(pl.col("ID1").lookup(T1, right_on="ID").alias("T1_ID1_match"))
    # 匹配ID2到T1
    .with_columns(pl.col("ID2").lookup(T1, right_on="ID").alias("T1_ID2_match"))
    # 匹配ID1到T2
    .with_columns(pl.col("ID1").lookup(T2, right_on="ID").alias("T2_ID1_match"))
    # 匹配ID2到T2
    .with_columns(pl.col("ID2").lookup(T2, right_on="ID").alias("T2_ID2_match"))
    # 展开struct列，提取目标值并重命名
    .unnest("T1_ID1_match")
    .rename({"Values": "ID1_Values", "Values II": "ID1_Values II"})
    .drop("ID")
    .unnest("T1_ID2_match")
    .rename({"Values": "ID2_Values", "Values II": "ID2_Values II"})
    .drop("ID")
    .unnest("T2_ID1_match")
    .rename({"Values": "ID1_T2_Values"})
    .drop("ID")
    .unnest("T2_ID2_match")
    .rename({"Values": "ID2_T2_Values"})
    .drop("ID")
)

print(result)

输出结果与方法1完全一致，且lookup是Polars优化过的向量操作，处理大表时性能远高于手动字典映射或多次join。

关于when的补充说明 #

如果一定要用when实现，需要结合map_elements做逐行匹配，但这种方法效率极低（每行都要执行一次过滤操作），仅适合演示：

# 不推荐用于大表，仅作语法演示
result = df.with_columns(
    ID1_Values = pl.when(pl.col("ID1").is_in(T1["ID"]))
                  .then(pl.col("ID1").map_elements(lambda x: T1.filter(pl.col("ID") == x)["Values"].item()))
                  .otherwise(None)
)

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