python-polars is there a np.where equivalent?

Ross 0
Viewed 69

Polars is there a np.where equivalent? trying to replicate the following code in polars. If the value is above a certain threshold column called Is_Acceptable is 1 or if it is below it is 0

import pandas as pd
import numpy as np 

df = pd.DataFrame({"fruit":["orange","apple","mango","kiwi"], "value":[1,0.8,0.7,1.2]})
df["Is_Acceptable?"] = np.where(df["value"].lt(0.9), 1, 0)
print(df)
python-polars
2 Answers

Yes, there is pl.when().then().otherwise() expression

import polars as pl
from polars import col

df = pl.DataFrame({
    "fruit": ["orange","apple","mango","kiwi"],
    "value": [1, 0.8, 0.7, 1.2]
})

df = df.with_column(
    pl.when(col('value') < 0.9).then(1).otherwise(0).alias('Is_Acceptable?')
)
print(df)

┌────────┬───────┬────────────────┐
│ fruit  ┆ value ┆ Is_Acceptable? │
│ ---    ┆ ---   ┆ ---            │
│ str    ┆ f64   ┆ i64            │
╞════════╪═══════╪════════════════╡
│ orange ┆ 1.0   ┆ 0              │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ apple  ┆ 0.8   ┆ 1              │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ mango  ┆ 0.7   ┆ 1              │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ kiwi   ┆ 1.2   ┆ 0              │
└────────┴───────┴────────────────┘
glebcom0

The when/then/otherwise expression is a good general-purpose answer. However, in this case, one shortcut is to simply create a boolean expression.

(
    df
    .with_column(
        (pl.col('value') < 0.9).alias('Is_Acceptable')
    )
)

shape: (4, 3)
┌────────┬───────┬───────────────┐
│ fruit  ┆ value ┆ Is_Acceptable │
│ ---    ┆ ---   ┆ ---           │
│ str    ┆ f64   ┆ bool          │
╞════════╪═══════╪═══════════════╡
│ orange ┆ 1.0   ┆ false         │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ apple  ┆ 0.8   ┆ true          │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ mango  ┆ 0.7   ┆ true          │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ kiwi   ┆ 1.2   ┆ false         │
└────────┴───────┴───────────────┘

In numeric computations, False will be upcast to 0, and True will be upcast to 1. Or, if you prefer, you can upcast them explicitly to a different type.

(
    df
    .with_column(
        (pl.col('value') < 0.9).cast(pl.Int64).alias('Is_Acceptable')
    )
)

shape: (4, 3)
┌────────┬───────┬───────────────┐
│ fruit  ┆ value ┆ Is_Acceptable │
│ ---    ┆ ---   ┆ ---           │
│ str    ┆ f64   ┆ i64           │
╞════════╪═══════╪═══════════════╡
│ orange ┆ 1.0   ┆ 0             │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ apple  ┆ 0.8   ┆ 1             │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ mango  ┆ 0.7   ┆ 1             │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ kiwi   ┆ 1.2   ┆ 0             │
└────────┴───────┴───────────────┘
cbilot0
Related