如何在极坐标中的多个条件下使用“when”、“then”和“otherwise”？

我有一个包含三列的数据集。将检查 A 列中的字符串。如果字符串匹配foo or spam，同一行中其他两列的值L and G应该改为XX。为此我尝试了以下方法。

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_column(
    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")))
)

然而，这是行不通的。有人可以帮我弄这个吗？

要将多列设置为相同的值，您可以使用：

df.with_columns(
   pl.when(pl.col("A").is_in(["foo", "spam"]))
     .then(pl.lit("XX"))
     .otherwise(pl.col("L", "G"))
     .keep_name()
)

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

.is_in()可以使用而不是多个== x | == y chains.

要使用不同的值一次更新多个列，您可以使用.map()和一本字典：

df.with_columns(
   pl.when(pl.col("A").is_in(["foo", "spam"]))
     .then(pl.col("L", "G").map(
        lambda col: {
           "L": "XX",
           "G": "YY",
        }.get(col.name)))
     .otherwise(pl.col("L", "G"))
)

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

使用时.map整个列都被传递（作为pl.Series）意味着您也可以根据需要操纵它：

df.with_columns(
   pl.when(pl.col("A").is_in(["foo", "spam"]))
     .then(pl.col("L", "G").map(
        lambda col: {
           "L": col.str.to_lowercase(),
           "G": "prefix_" + col
        }.get(col.name)))
     .otherwise(pl.col("L", "G"))
)

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