如何在极坐标条件下生成多个

Question

我有一本字典，其中字符串作为键，极坐标表达式作为值。

我怎样才能以简洁的方式做这样的事情：

df = df.with_columns(
    pl.when(condition_1)
    .then(pl.lit(key_1))
    .when(pl.lit(condition_2))
    .then(pl.lit(key_2))
    ...
    .otherwise(None)
    .alias("new_column")
)

Answer 1

考虑以下示例数据。

import polars as pl

df = pl.DataFrame({
    "num": list(range(6)),
})

shape: (6, 1)
┌─────┐
│ num │
│ --- │
│ i64 │
╞═════╡
│ 0   │
│ 1   │
│ 2   │
│ 3   │
│ 4   │
│ 5   │
└─────┘

一般来说，

pl.when().then().otherwise()

构造可以嵌套以获得 switch 语句的效果，您似乎在问题中概述了这一点。

df.with_columns(
    pl.when(
        pl.col("num") < 2
    ).then(
        pl.lit("small")
    ).otherwise(
        pl.when(
            pl.col("num") > 3
        ).then(
            pl.lit("large")
        ).otherwise(
            pl.lit("medium")
        )
    )
)

shape: (6, 2)
┌─────┬─────────┐
│ num ┆ literal │
│ --- ┆ ---     │
│ i64 ┆ str     │
╞═════╪═════════╡
│ 0   ┆ small   │
│ 1   ┆ small   │
│ 2   ┆ medium  │
│ 3   ┆ medium  │
│ 4   ┆ large   │
│ 5   ┆ large   │
└─────┴─────────┘

如果嵌套了许多条件，这可能会很乏味。在这种情况下，如果不满足

pl.coalesce

中的条件，则

pl.when().then()

构造将求值为 null，因此

pl.when()

可能会有所帮助。

df.with_columns(
    pl.coalesce(
        pl.when(pl.col("num") < 2).then(pl.lit("small")),
        pl.when(pl.col("num") > 3).then(pl.lit("large")),
        pl.lit("medium")
    )
)

shape: (6, 2)
┌─────┬─────────┐
│ num ┆ literal │
│ --- ┆ ---     │
│ i64 ┆ str     │
╞═════╪═════════╡
│ 0   ┆ small   │
│ 1   ┆ small   │
│ 2   ┆ medium  │
│ 3   ┆ medium  │
│ 4   ┆ large   │
│ 5   ┆ large   │
└─────┴─────────┘

如果您有一个字典，其中值作为键，条件作为值，则可以按如下方式使用。

d = {
    "small": pl.col("num") < 2,
    "medium": pl.col("num") < 4,
    "large": pl.col("num") >= 4,
}

df.with_columns(
    pl.coalesce(
        pl.when(cond).then(pl.lit(val)) for val, cond in d.items()
    )
)