Conditional expressions in Python

when/then/otherwise is the vectorized if, and it chains.

import polars as pl

frame = pl.DataFrame({"tpm": [104, 88, 70, 96]})

graded = frame.with_columns(
    pl.when(pl.col("tpm") >= 100).then(pl.lit(3))
      .when(pl.col("tpm") >= 90).then(pl.lit(2))
      .otherwise(pl.lit(1))
      .alias("stars"),
    pl.when(pl.col("tpm") > pl.col("tpm").mean())
      .then(pl.lit("above"))
      .otherwise(pl.lit("below"))
      .alias("vs_mean"),
)
print(graded)

How it works

Each when adds a branch, tested in order.
otherwise is the fallback and is required for full coverage.
The whole chain is one expression, so it runs column-wise.

Keywords and builtins used here

as
print

The run, in numbers

Lines: 15
Characters to type: 359
Tokens: 164
Three-star pace: 105 tpm

At the three-star pace of 105 tokens a minute, this run takes about 94 seconds.

Type this snippet

Step 4 of 4 in Expressions, step 8 of 32 in Polars.

← Previous Next →