我有一本字典,其中字符串作为键,极坐标表达式作为值。
我怎样才能以简洁的方式做这样的事情:
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")
)
考虑以下示例数据。
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()
)
)