How to add a conditional computed col to polars dataframe in rust?

Question

df1 has ['a', 'b', 'c'] 3 cols, I want to get a df2 with 4 cols of ['a', 'b', 'c', 'd']. And d is computed like this:

if a>5 {
  d = b + c 
} else if a<-5 {
  d = c - b + a
}  else {
  d = 3.0 * a
}

How can I do this with polars in rust? maybe for both eager and lazy.

Answer 1

You can use a when -> then -> when -> then -> otherwise expression. Note that you can extend the when, then indefinitely, just like an else if branch.

Below is an example:

use polars::df;
use polars::prelude::*;

fn main() -> Result<()> {
    let df = df![
        "a" => [2, 9, 2, 5],
        "b" => [1, 2, 3, 4],
        "c" => [4, 4, 8, 4],
    ]?;

    let out = df
        .lazy()
        .select([
            col("*"),
            when(col("a").gt(lit(5)))
                .then(col("b") + col("c"))
                .when(col("a").lt(lit(5)))
                .then(col("c") - col("b") + col("a"))
                .otherwise(lit(3) * col("a"))
                .alias("d"),
        ])
        .collect()?;

    println!("{}", out);

    Ok(())
}

This outputs:

shape: (4, 4)
┌─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   │
│ --- ┆ --- ┆ --- ┆ --- │
│ i32 ┆ i32 ┆ i32 ┆ i32 │
╞═════╪═════╪═════╪═════╡
│ 2   ┆ 1   ┆ 4   ┆ 5   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 9   ┆ 2   ┆ 4   ┆ 6   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2   ┆ 3   ┆ 8   ┆ 7   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 5   ┆ 4   ┆ 4   ┆ 15  │
└─────┴─────┴─────┴─────┘