PyPolars：无需循环，根据另一列值获取指定列数据

你的原循环写法有个隐藏问题：每次调用with_columns生成new_col时，都会覆盖之前的结果，最终只有最后一组映射（比如你的例子里D→col_3）会生效，前面的条件逻辑都被冲掉了。下面用Polars原生方法完美解决这个问题，分两种常用方案：

方案一：动态列选择（最简洁） #

利用pl.col()支持表达式参数的特性，先通过map_dict把col_val的值映射到目标列名，再动态选取对应列的数值，一行代码搞定：

import polars as pl

def do_that_mapping(lf: pl.LazyFrame) -> pl.LazyFrame:
    my_map = {
        "A": "col_1",
        "B": "col_2",
        "C": "col_2",
        "D": "col_3",
        # 更多映射规则
    }
    return lf.with_columns(
        pl.col(pl.col("col_val").map_dict(my_map)).alias("new_col")
    )

# 测试
x = pl.LazyFrame(
    data={
        "col_val": ["A", "B", "C", "D"],
        "col_1": [22, 1, 54, 82],
        "col_2": [1, 32, 7, 8],
        "col_3": [4, 6, 90, 3],
    },
    schema={
        "col_val": pl.String,
        "col_1": pl.Int16,
        "col_2": pl.Int16,
        "col_3": pl.Int16,
    },
)
x_with_new_col = x.pipe(do_that_mapping).collect()
print(x_with_new_col)

输出结果：

shape: (4, 4)
┌─────────┬──────┬──────┬──────┬────────┐
│ col_val ┆ col_1┆ col_2┆ col_3┆ new_col│
│ ---     ┆ ---  ┆ ---  ┆ ---  ┆ ---    │
│ str     ┆ i16  ┆ i16  ┆ i16  ┆ i16    │
╞═════════╪══════╪══════╪══════╪════════╡
│ A       ┆ 22   ┆ 1    ┆ 4    ┆ 22     │
│ B       ┆ 1    ┆ 32   ┆ 6    ┆ 32     │
│ C       ┆ 54   ┆ 7    ┆ 90   ┆ 7      │
│ D       ┆ 82   ┆ 8    ┆ 3    ┆ 3      │
└─────────┴──────┴──────┴──────┴────────┘

方案二：链式when/then（适合复杂条件扩展） #

如果需要添加更多自定义逻辑（比如默认值、多条件判断），可以用functools.reduce构建链式的when/then表达式，避免显式循环：

import polars as pl
from functools import reduce

def do_that_mapping(lf: pl.LazyFrame) -> pl.LazyFrame:
    my_map = {
        "A": "col_1",
        "B": "col_2",
        "C": "col_2",
        "D": "col_3",
        # 更多映射规则
    }
    # 构建初始空条件，链式拼接所有when/then
    condition = reduce(
        lambda expr, (k, v): expr.when(pl.col("col_val") == k).then(pl.col(v)),
        my_map.items(),
        pl.when(False)  # 初始条件为False，确保第一个when生效
    )
    # 可添加otherwise处理未匹配的情况，比如 .otherwise(pl.lit(0))
    return lf.with_columns(condition.alias("new_col"))

这个方案和方案一效果一致，但更灵活，能轻松扩展额外的条件逻辑。

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