如何在Polars的when-then-otherwise表达式中批量更新多列？

报错原因 #

你遇到的ComputeError是因为Polars不允许在单个when-then-otherwise表达式中直接返回多列并通过name.map批量重命名替换，with_columns无法处理这种多列扩展逻辑。

解决方案 #

以下两种方式均可实现需求：

方法1：为目标列分别编写独立逻辑

先定义复用的判断条件，再针对b1和b2各自实现when-then-otherwise逻辑：

import polars as pl

df = pl.DataFrame({
    'a1': [1,1,1], 
    'a2': [2,2,None], 
    'b1': [3,None,None], 
    'b2': [4,None,None]
})

# 定义判断条件
condition = (
    pl.col('a1').is_not_null() &
    pl.col('a2').is_not_null() &
    pl.col('b1').is_null() &
    pl.col('b2').is_null()
)

# 对b1、b2分别赋值
result = df.with_columns(
    pl.when(condition).then(pl.col('a1')).otherwise(pl.col('b1')).alias('b1'),
    pl.when(condition).then(pl.col('a2')).otherwise(pl.col('b2')).alias('b2')
)

print(result)

方法2：使用struct批量处理列

通过pl.struct打包需要处理的列，完成逻辑后再展开，减少重复代码：

import polars as pl

df = pl.DataFrame({
    'a1': [1,1,1], 
    'a2': [2,2,None], 
    'b1': [3,None,None], 
    'b2': [4,None,None]
})

condition = (
    pl.col('a1').is_not_null() &
    pl.col('a2').is_not_null() &
    pl.col('b1').is_null() &
    pl.col('b2').is_null()
)

result = df.with_columns(
    pl.when(condition)
    # 满足条件时用a1/a2生成对应b1/b2的struct
    .then(pl.struct(b1=pl.col('a1'), b2=pl.col('a2')))
    # 不满足条件时保留原有b1/b2的struct
    .otherwise(pl.struct(b1=pl.col('b1'), b2=pl.col('b2')))
    .alias('temp_struct')
).unnest('temp_struct')  # 展开struct为列

print(result)

验证结果 #

两种方法都会输出符合预期的DataFrame：

┌─────┬──────┬──────┬──────┐
│ a1  ┆ a2   ┆ b1   ┆ b2   │
│ --- ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╡
│ 1   ┆ 2    ┆ 3    ┆ 4    │
│ 1   ┆ 2    ┆ 1    ┆ 2    │
│ 1   ┆ null ┆ null ┆ null │
└─────┴──────┴──────┴──────┘

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