Как сделать несколько при условии на полярах

Рассмотрим следующий пример данных.

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() могут быть вложены, чтобы получить эффект оператора переключения, который вы, кажется, обрисовали в общих чертах в своем вопросе.

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() оценивается как нулевое, если условие в 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()
    )
)