基于单一条件批量重分配Polars多列的更优雅实现方法

问题背景 #

我正在学习Polars，经常碰到需要基于单一条件对多列进行值重分配的场景，对应的伪代码逻辑如下：

if x=2 then do:
  x='ham'
  y='eggs'
  z='fruit'
end

目前我用以下Python-Polars代码实现了需求，但觉得写法不够优雅：

df = pl.DataFrame(
    {
        "a": [1,2,3],
        "x": ['foo','bar','baz'],
        "y": ['bar','foo','baz'],
        "z": ['baz','foo','bar']
    }
)
df.with_columns(
    pl.when(pl.col('a')==2).then(pl.lit('ham')).otherwise(pl.col('x')).alias('x'),
    pl.when(pl.col('a')==2).then(pl.lit('eggs')).otherwise(pl.col('y')).alias('y'),
    pl.when(pl.col('a')==2).then(pl.lit('fruit')).otherwise(pl.col('z')).alias('z'),
    
)

更优实现方案 #

方案1：字典批量生成列表达式 #

把需要更新的列与对应值存入字典，通过列表推导式批量生成when/then/otherwise表达式，避免重复编写相同的条件判断：

update_map = {"x": "ham", "y": "eggs", "z": "fruit"}
condition = pl.col("a") == 2

df.with_columns(
    [pl.when(condition).then(pl.lit(val)).otherwise(pl.col(col)).alias(col) 
     for col, val in update_map.items()]
)

后续要新增更新列时，只需要扩展update_map字典即可，代码更简洁易维护。

方案2：Struct打包+批量处理 #

将目标列打包成struct，统一处理后再展开，把多列更新逻辑聚合在一起：

condition = pl.col("a") == 2

df.with_columns(
    pl.struct("x", "y", "z")
      .when(condition)
      .then(pl.struct(x=pl.lit("ham"), y=pl.lit("eggs"), z=pl.lit("fruit")))
      .unnest()
)

这种方式减少了重复的条件判断，逻辑更集中，适合后续更新逻辑可能变得复杂的场景。

方案3：直接使用replace方法 #

Polars的replace方法支持直接指定条件和替换的键值对，语法最直观：

df.replace(
    where=pl.col("a") == 2,
    replacement={"x": "ham", "y": "eggs", "z": "fruit"}
)

这个方法专门针对条件替换场景，代码量最少，可读性最高。

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