TopolearnBucketing, flags and business rules are everywhere, and when/then is the vectorized way to write them.
Python if statements can't decide per row inside an expression. Polars uses pl.when(condition).then(value).otherwise(default) instead, which is like SQL's CASE WHEN.
Chain more branches with extra .when(...).then(...) pairs; the first condition that is true wins, top to bottom. If you leave out .otherwise(...), rows that match no branch become null.
pl.when(pl.col("score") >= 90).then(pl.lit("A"))
.when(pl.col("score") >= 75).then(pl.lit("B"))
.otherwise(pl.lit("C"))Values in then and otherwise can be literals or other expressions, so you can pick between columns: pl.when(pl.col("discount").is_not_null()).then(pl.col("discount")).otherwise(pl.col("price")). For that particular "first non-null" pattern, pl.coalesce("discount", "price") is shorter.
A string passed to then("x") is read as a column name, not a literal. Wrap text in pl.lit("x") when you mean the string itself.
All branches are computed for all rows and then combined, which is fast because it's vectorized. The result type is the common supertype of the branches, so mixing ints and floats gives floats.
Example
import polars as pl
df = pl.DataFrame({"name": ["Ana", "Ben", "Cy", "Di"],
"score": [95, 80, 60, None]})
graded = df.with_columns(
grade=pl.when(pl.col("score") >= 90).then(pl.lit("A"))
.when(pl.col("score") >= 75).then(pl.lit("B"))
.when(pl.col("score").is_not_null()).then(pl.lit("C"))
# no otherwise: rows matching nothing become null
)
print(graded["grade"].to_list()) # ['A', 'B', 'C', None]
bonus = pl.when(pl.col("grade") == "A").then(10).otherwise(0)
print(graded.select(bonus.alias("bonus"))["bonus"].to_list()) # [10, 0, 0, 0]Remember
pl.when(...).then(...).otherwise(...)is a vectorized CASE WHEN- The first true branch wins; no otherwise means null
- Wrap string values in
pl.lit(); bare strings are column names
Watch out for
then("A")when you meant the text A: Polars looks for a column called A.