如何在Polars中结合when-then实现多列条件计算

需求说明 #

针对包含字符串列A、L、G的Polars DataFrame：

• L和G列格式固定为「单个字母+两位数字」
• 当A列值为"foo"或"spam"时：
  • 将L列修改为"XX"
  • 将G列修改为「原G列的字母部分」拼接「原L列的数字部分」
• 其余行保持原列值不变

代码实现 #

import polars as pl

# 构造示例DataFrame
df = pl.DataFrame(
    {
        "A": ["foo", "ham", "spam", "egg"],
        "L": ["A54", "A12", "B84", "C12"],
        "G": ["X34", "C84", "G96", "L60"],
    }
)

# 执行条件修改
result_df = df.with_columns(
    # 修改L列
    pl.when(pl.col("A").is_in(["foo", "spam"]))
      .then("XX")
      .otherwise(pl.col("L"))
      .alias("L"),
    # 修改G列：取G的首字母 + L的后两位数字
    pl.when(pl.col("A").is_in(["foo", "spam"]))
      .then(pl.col("G").str.slice(0, 1) + pl.col("L").str.slice(1, 2))
      .otherwise(pl.col("G"))
      .alias("G")
)

print(result_df)

执行结果 #

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

代码说明 #

• 用pl.col("A").is_in(["foo", "spam"])精准匹配目标行
• 对于G列的字符串拆分：
  • str.slice(0, 1)提取G列的第一个字符（字母部分）
  • str.slice(1, 2)提取L列从第2位开始的2个字符（数字部分）
• 通过with_columns同时完成两列的条件更新，保证处理效率

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