How to use with_column method to create a calculated column in Polars Rust?

Question

I was trying to create a new computed column based on existing column in polars rust DataFrame. There is a pyspark like with_column method available for that. But in the api documentation there is no example. Here is a example dataframe:

use polars::prelude::*;

fn example() {
    let df = df!["foo" => ["A", "A", "B", "B", "C"],
        "val1" => [1, 2, 2, 4, 2],
        "val2" => [1, 2, 2, 4, 2],
        ].unwrap();
    // newcolumn ration = val1/val2
    // df.with_column(...)
    println!("{}", df);

fn main{
    example()
}

I want to create a ration column which will calculate the ration between val1 and val 2 but there is no example available in the API documentation. Also there is another issue. The with column method might also need the col type to wrap the columns like pyspark but polars::prelute::* does not brings the col type into scope. Or may be some features needed to be enabled in the cargo file.

I am using latest version of Polars 0.22.8.

Does any one knows how to do it.

Answer 1

Your initial idea works with the lazy API:

# Cargo.toml
# ...
[dependencies]
polars = { version = "0.22.8", features = ["lazy"] }

// src/main.rs
use polars::prelude::*;

fn example() -> DataFrame {
    let df = df!["foo" => ["A", "A", "B", "B", "C"],
    "val1" => [1, 2, 2, 4, 2],
    "val2" => [1, 2, 2, 4, 2],
    ]
    .unwrap();

    df.lazy()
        .with_column((col("val1") / col("val2")).alias("ration"))
        .collect()
        .unwrap()
}

fn main() {
    let df = example();
    println!("{:?}", df);
}

Output:

shape: (5, 4)
┌─────┬──────┬──────┬────────┐
│ foo ┆ val1 ┆ val2 ┆ ration │
│ --- ┆ ---  ┆ ---  ┆ ---    │
│ str ┆ i32  ┆ i32  ┆ i32    │
╞═════╪══════╪══════╪════════╡
│ A   ┆ 1    ┆ 1    ┆ 1      │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ A   ┆ 2    ┆ 2    ┆ 1      │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ B   ┆ 2    ┆ 2    ┆ 1      │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ B   ┆ 4    ┆ 4    ┆ 1      │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ C   ┆ 2    ┆ 2    ┆ 1      │
└─────┴──────┴──────┴────────┘