Polars：在when/then表达式中用列值动态引用其他列

问题场景 #

需要生成新列res，规则如下：

• 当frequency列非空且非NaN时，直接取frequency的值
• 否则根据frequency_ref列指定的列名，从对应列取值
• 兜底值为0.0

示例数据：

import polars as pl

df = pl.DataFrame(
    {
        "frequency": [0.5, None, None, None],
        "frequency_ref": ["a", "z", "a", "a"],
        "a": [1, 2, 3, 4],
        "z": [5, 6, 7, 8],
    }
)

期望结果：

res = pl.DataFrame(
    {
        "frequency": [0.5, None, None, None],
        "frequency_ref": ["a", "z", "a", "a"],
        "a": [1, 2, 3, 4],
        "z": [5, 6, 7, 8],
        "res": [0.5, 6, 3, 4]
    }
)

错误原因 #

你尝试的pl.col(pl.col("frequency_ref"))写法不合法：pl.col()要求传入字符串（列名）或DataType类型，不能直接传入表达式（Expr对象），因此触发TypeError: invalid input for "col". Expected "str" or "DataType", got 'Expr'.

最优解决方案 #

方法一：利用结构体（Struct）动态取值（推荐，性能更优） #

将需要动态引用的列打包成结构体，通过frequency_ref的值作为键提取对应字段：

# 定义所有可能被引用的列集合
target_cols = ["a", "z"]

df = df.with_columns(
    pl.when(pl.col("frequency").is_not_null() & pl.col("frequency").is_not_nan())
    .then(pl.col("frequency"))
    # 可选：先验证引用列是否在目标集合中，避免无效列名报错
    .when(pl.col("frequency_ref").is_in(target_cols))
    .then(pl.struct(target_cols).struct.field(pl.col("frequency_ref")))
    .otherwise(0.0)
    .alias("res")
)

优势：

• 无需硬编码每个列的when/then分支
• 列名变更时仅需修改target_cols集合，扩展性强
• 基于Polars向量化操作，大数据集下性能更优

方法二：利用pl.row动态索引列 #

通过pl.row接收表达式返回的列名字符串，直接取对应列的值：

df = df.with_columns(
    pl.when(pl.col("frequency").is_not_null() & pl.col("frequency").is_not_nan())
    .then(pl.col("frequency"))
    .then(pl.row(pl.col("frequency_ref")))
    .otherwise(0.0)
    .alias("res")
)

优势：代码更简洁，无需提前定义目标列集合；但pl.row是按行处理操作，大数据集下性能略逊于结构体方法。

对比硬编码方案 #

硬编码逐个列的when/then分支扩展性极差，列名新增或修改时需同步修改代码；上述两种方法仅需维护目标列集合或自动适配，大幅降低维护成本。

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