如何在Polars中基于外部条件对指定多列进行条件赋值？

问题背景 #

现有如下Polars DataFrame：

import polars as pl

df = pl.DataFrame({'a':[0,1,2,3,4],
                   'b':[10,20,30,40,50],
                   'c':[100,200,300,400,500]})

输出结果：

┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 0   ┆ 10  ┆ 100 │
│ 1   ┆ 20  ┆ 200 │
│ 2   ┆ 30  ┆ 300 │
│ 3   ┆ 40  ┆ 400 │
│ 4   ┆ 50  ┆ 500 │
└─────┴─────┴─────┘

同时有外部条件列表和目标列列表：

condition = [True, False, False, False, True]
columns = ['a', 'c']

在Pandas中可通过df.loc[condition, columns] = None实现基于该条件对指定列赋值，需要在Polars中实现相同效果，得到如下结果：

┌──────┬─────┬──────┐
│ a    ┆ b   ┆ c    │
│ ---  ┆ --- ┆ ---  │
│ i64  ┆ i64 ┆ i64  │
╞══════╪═════╪══════╡
│ null ┆ 10  ┆ null │
│ 1    ┆ 20  ┆ 200  │
│ 2    ┆ 30  ┆ 300  │
│ 3    ┆ 40  ┆ 400  │
│ null ┆ 50  ┆ null │
└──────┴─────┴──────┘

解决方案 #

在Polars中，可以通过with_columns结合when/then/otherwise表达式实现需求，代码如下：

result_df = df.with_columns(
    [pl.when(pl.Series(condition)).then(None).otherwise(pl.col(col)).alias(col) for col in columns]
)

print(result_df)

代码说明 #

1. 条件适配：用pl.Series(condition)将外部Python列表转为Polars Series，确保能和DataFrame的行一一对应；
2. 列处理逻辑：遍历目标列列表，对每个列应用条件判断：当对应行的condition为True时赋值None（Polars中显示为null），否则保留原列值；
3. 列名保持：通过alias(col)确保处理后的列名与原列一致；
4. 类型自动转换：原始的非可空整数列（i64）会自动转为可空整数类型（Int64），以支持null值存储，这与Pandas的行为一致。

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