Conditional Logic or If-Else in Polars

One of the most common patterns in data is the need for conditional logic or if-else statements. You may want to assign a specific value when a certain condition is met. This is easily done in Pandas, using numpy.where() or pandas.where(). But how about Polars? The short answer is yes, you can do the same in Polars. You’d use polars.when().

You can find my Python script used in this post in my GitHub repo.

Implement Conditional Logic in Polars

As mentioned above, you use polars.when() to implement conditional logic or if-else statements in Polars.

Let’s say I have a dataframe that contains numbers from 1 to 5.

import polars as pl

df = pl.LazyFrame(
    {'Numbers': [1,2,3,4,5]}
)

print(df.collect().head())
'''
output:
shape: (5, 1)
┌─────────┐
│ Numbers │
│ ---     │
│ i64     │
╞═════════╡
│ 1       │
│ 2       │
│ 3       │
│ 4       │
│ 5       │
└─────────┘
'''

Now I add the portion that adds a new column assigning values depending on the value of “Number” column. Use .when() to define your condition and use .otherwise() to assigning “else” values.

df = (
    df
    .with_columns(
        pl.when(pl.col('Numbers')==1)
        .then('Best')
        .when(pl.col('Numbers')==2)
        .then('Second Best')
        .otherwise('Not Good')
        .alias('Rank')
    )
    
)

print(df.collect().head())
'''
output:
shape: (5, 2)
┌─────────┬─────────────┐
│ Numbers ┆ Rank        │
│ ---     ┆ ---         │
│ i64     ┆ str         │
╞═════════╪═════════════╡
│ 1       ┆ Best        │
│ 2       ┆ Second Best │
│ 3       ┆ Not Good    │
│ 4       ┆ Not Good    │
│ 5       ┆ Not Good    │
└─────────┴─────────────┘
'''

Summary

Your code can be a little lengthy when using polars.when(), but it’s much more readable than pandas.where() in my opinion as polars.when() is more descriptive and easy to read the coder’s intention.

GitHub repo

Reference

https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.when.html