如何用Polars DataFrame高效添加inverse列（匹配反向value值）

针对给定的Polars DataFrame，要添加inverse列（对应value(y,x)，不存在则为0），最简洁高效的方式是使用自左连接：

import polars as pl

df = pl.DataFrame({
    'x':['a','a','b','b'],
    'y':['b','c','c','a'],
    'value':[3,5,1,4]
})

result = df.join(
    # 重命名列，将原表的(x,y)转换为用于匹配的(y,x)
    df.rename({'x': 'match_y', 'y': 'match_x', 'value': 'inverse'}),
    left_on=['x', 'y'],
    right_on=['match_x', 'match_y'],
    how='left'
).fill_null(0).drop(['match_x', 'match_y'])

print(result)

执行后输出：

shape: (4, 4)
┌─────┬─────┬───────┬─────────┐
│ x   ┆ y   ┆ value ┆ inverse │
│ --- ┆ --- ┆ ---   ┆ ---     │
│ str ┆ str ┆ i64   ┆ i64     │
╞═════╪═════╪═══════╪═════════╡
│ a   ┆ b   ┆ 3     ┆ 4       │
│ a   ┆ c   ┆ 5     ┆ 0       │
│ b   ┆ c   ┆ 1     ┆ 0       │
│ b   ┆ a   ┆ 4     ┆ 3       │
└─────┴─────┴───────┴─────────┘

替代方案（适合小数据集） #

如果数据集较小，也可以用字典映射+表达式的方式快速实现：

# 构建(y,x)到value的映射字典
inv_map = df.select(pl.struct("y", "x").alias("key"), "value").to_dicts()
inv_dict = {tuple(d['key'].values()): d['value'] for d in inv_map}

# 添加inverse列
result = df.with_columns(
    pl.struct("x", "y").map_elements(
        lambda s: inv_dict.get(tuple(s.values()), 0),
        return_dtype=pl.Int64
    ).alias("inverse")
)

性能说明 #

• 自连接方法是Polars原生优化的操作，在大数据量下性能远优于Python层面的map_elements，是优先推荐的方案。
• 字典映射方法代码更直观，但仅适合小数据集，因为map_elements会触发Python UDF，无法利用Polars的矢量化加速。

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