如何在Python Polars中基于列条件为多列赋相同值

原始数据 #

import polars as pl

df = pl.DataFrame({
   "A": ["a", "b", "a", "c"], 
   "B": [23,45,60,50], 
   "C": [11, 26, 63, 84]
})

原始数据展示：

┌─────┬─────┬─────┐
│ A   ┆ B   ┆ C   │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ a   ┆ 23  ┆ 11  │
│ b   ┆ 45  ┆ 26  │
│ a   ┆ 60  ┆ 63  │
│ c   ┆ 50  ┆ 84  │
└─────┴─────┴─────┘

需求 #

当列A的值为"a"时，将列B和C的值设置为0，期望输出：

┌─────┬─────┬─────┐
│ A   ┆ B   ┆ C   │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ a   ┆ 0   ┆ 0   │
│ b   ┆ 45  ┆ 26  │
│ a   ┆ 0   ┆ 0   │
│ c   ┆ 50  ┆ 84  │
└─────┴─────┴─────┘

解法 #

方法1：逐列使用when/then/otherwise #

通过pl.when()判断条件，满足时设为0，否则保留原值：

result = df.with_columns(
    pl.when(pl.col("A") == "a").then(0).otherwise(pl.col("B")).alias("B"),
    pl.when(pl.col("A") == "a").then(0).otherwise(pl.col("C")).alias("C")
)
print(result)

方法2：批量处理多列（更简洁高效） #

对指定列列表批量应用条件逻辑，避免重复代码：

result = df.with_columns(
    pl.col(["B", "C"]).when(pl.col("A") == "a").then(0).keep_name()
)
print(result)

两种方法都能得到期望的输出，方法2在需要处理多列时更简洁，适合扩展。

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