如何基于Polars数据框中列的取值，通过表达式字典映射生成新列？

当然可以实现啦！你的思路方向是对的，但直接用apply处理表达式字典并不是Polars里最优的做法——毕竟apply是逐行处理，没法发挥Polars向量化操作的性能优势，而且还容易出问题。咱们换个更贴合Polars风格的方式来实现，既高效又靠谱。

方法一：手动构建when/then条件链 #

我们可以利用Polars的条件表达式when/then，把字典里的每个键值对转换成对应的条件逻辑，代码如下：

import polars as pl

df = pl.DataFrame({
    "col1": ["a", "b", "a"],
    "x": [1,2,3],
    "y": [2,2,5]
})

expr_dict = {
    "a": pl.col("x") * pl.col("y"),
    "b": pl.col("x"),
}

# 初始化条件表达式
condition_expr = None
for key, expr in expr_dict.items():
    if condition_expr is None:
        # 第一个条件直接用when创建
        condition_expr = pl.when(pl.col("col1") == key).then(expr)
    else:
        # 后续条件用when链式拼接
        condition_expr = condition_expr.when(pl.col("col1") == key).then(expr)

# 处理col1不在字典中的情况，这里设为None，你可以改成其他默认值
condition_expr = condition_expr.otherwise(None)

# 生成新列r
result_df = df.with_columns(r=condition_expr)
print(result_df)

运行这段代码后，你就能得到想要的结果：

shape: (3, 4)
┌──────┬─────┬─────┬─────┐
│ col1 ┆ x   ┆ y   ┆ r   │
│ ---  ┆ --- ┆ --- ┆ --- │
│ str  ┆ i64 ┆ i64 ┆ i64 │
╞══════╪═════╪═════╪═════╡
│ a    ┆ 1   ┆ 2   ┆ 2   │
│ b    ┆ 2   ┆ 2   ┆ 2   │
│ a    ┆ 3   ┆ 5   ┆ 15  │
└──────┴─────┴─────┴─────┘

方法二：用reduce简化条件链写法 #

如果你的表达式字典比较大，手动循环拼接有点繁琐，还可以用functools.reduce来简化代码，让逻辑更紧凑：

import polars as pl
from functools import reduce

df = pl.DataFrame({
    "col1": ["a", "b", "a"],
    "x": [1,2,3],
    "y": [2,2,5]
})

expr_dict = {
    "a": pl.col("x") * pl.col("y"),
    "b": pl.col("x"),
}

# 用reduce串联所有when/then条件
condition_expr = reduce(
    lambda acc, (key, expr): acc.when(pl.col("col1") == key).then(expr),
    expr_dict.items(),
    pl.when(False)  # 初始值设为永远不成立的条件
).otherwise(None)

result_df = df.with_columns(r=condition_expr)
print(result_df)

小提示 #

• 如果col1的所有可能取值都在expr_dict里，那otherwise(None)可以省略，但加上会让代码更健壮，避免遇到未知键时出现意外的空值或错误。
• 这两种方法都是基于Polars的向量化操作，性能比apply好很多，尤其是处理大规模数据集的时候，差距会非常明显。

备注：内容来源于stack exchange，提问作者Babak Fi Foo