如何在Polars中用when、then、otherwise处理多条件列修改

我有一个包含三列的数据集，需检查A列的字符串：当A列值匹配foo或spam时，将同行的L列和G列值改为XX。

数据集示例 #

import polars as pl

df = pl.DataFrame(
    {
        "A": ["foo", "ham", "spam", "egg",],
        "L": ["A54", "A12", "B84", "C12"],
        "G": ["X34", "C84", "G96", "L6",],
    }
)
print(df)

输出结果：

shape: (4, 3)
┌──────┬─────┬─────┐
│ A ┆ L ┆ G │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str │
╞══════╪═════╪═════╡
│ foo ┆ A54 ┆ X34 │
│ ham ┆ A12 ┆ C84 │
│ spam ┆ B84 ┆ G96 │
│ egg ┆ C12 ┆ L6 │
└──────┴─────┴─────┘

期望输出 #

shape: (4, 3)
┌──────┬─────┬─────┐
│ A ┆ L ┆ G │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str │
╞══════╪═════╪═════╡
│ foo ┆ XX ┆ XX │
│ ham ┆ A12 ┆ C84 │
│ spam ┆ XX ┆ XX │
│ egg ┆ C12 ┆ L6 │
└──────┴─────┴─────┘

尝试的错误代码 #

df = df.with_columns(
    pl.when((pl.col("A") == "foo") | (pl.col("A") == "spam"))
    .then((pl.col("L")= "XX") & (pl.col( "G")= "XX"))
    .otherwise((pl.col("L"))&(pl.col( "G")))
)

错误原因 #

1. Polars表达式中不能直接使用赋值语句pl.col("L")= "XX"，这不符合表达式语法规则。
2. 使用&连接列表达式是逻辑与运算，无法实现同时修改多列的需求。

正确实现方案 #

方案一：分别处理每一列 #

明确为L和G列各自定义条件逻辑：

df = df.with_columns(
    # 修改L列
    pl.when(pl.col("A").is_in(["foo", "spam"]))
      .then("XX")
      .otherwise(pl.col("L")).alias("L"),
    # 修改G列
    pl.when(pl.col("A").is_in(["foo", "spam"]))
      .then("XX")
      .otherwise(pl.col("G")).alias("G"),
)

方案二：批量处理多列 #

如果需要修改的列较多，用列表推导式批量生成表达式更高效：

target_cols = ["L", "G"]
df = df.with_columns(
    [
        pl.when(pl.col("A").is_in(["foo", "spam"]))
          .then("XX")
          .otherwise(pl.col(col)).alias(col)
        for col in target_cols
    ]
)

两种方案运行后都能得到期望的输出结果。

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