Ben Chuanlong Du's Blog

It is never too late to learn.

Map a Function to Multiple Columns in Polars in Rust

Things on this page are fragmentary and immature notes/thoughts of the author. Please read with your own judgement!

In [2]:
:dep polars = { version = "0.30", features = ["lazy", "round_series"] }
In [ ]:
fn my_filter_func(col1: &Series, col2: &Series, col2: &Series) -> ReturnType {
    let it = (0..n).map(|i| {
        let col1 = match col.get(i) {
            AnyValue::UInt64(val) => val,
            _ => panic!("Wrong type of col1!"),
        };
        // similar for col2 and col3
        // apply user-defined function to col1, col2 and col3
    }
    // convert it to a collection of the required type
}
In [ ]:
fn my_black_box_function(a: f32, b: f32) -> f32 {
    // do something
    a
}

fn apply_multiples(col_a: &Series, col_b: &Series) -> Float32Chunked {
    match (col_a.dtype(), col_b.dtype()) {
        (DataType::Float32, DataType::Float32) => {
            let a = col_a.f32().unwrap();
            let b = col_b.f32().unwrap();

            a.into_iter()
                .zip(b.into_iter())
                .map(|(opt_a, opt_b)| match (opt_a, opt_b) {
                    (Some(a), Some(b)) => Some(my_black_box_function(a, b)),
                    _ => None,
                })
                .collect()
        }
        _ => panic!("unpexptected dtypes"),
    }
}
In [ ]:
fn apply_multiples(lf: LazyFrame) -> Result<DataFrame> {
    df![
        "a" => [1.0, 2.0, 3.0],
        "b" => [3.0, 5.1, 0.3]
    ]?
    .lazy()
    .select([concat_lst(["col_a", "col_b"]).map(
        |s| {
            let ca = s.struct_()?;

            let b = ca.field_by_name("col_a")?;
            let a = ca.field_by_name("col_b")?;
            let a = a.f32()?;
            let b = b.f32()?;

            let out: Float32Chunked = a
                .into_iter()
                .zip(b.into_iter())
                .map(|(opt_a, opt_b)| match (opt_a, opt_b) {
                    (Some(a), Some(b)) => Some(my_black_box_function(a, b)),
                    _ => None,
                })
                .collect();

            Ok(out.into_series())
        },
        GetOutput::from_type(DataType::Float32),
    )])
    .collect()
}
In [3]:
use polars::{
    prelude::*,
    datatypes::DataType,
};
In [12]:
let df = df!(
    "column integers"  => &[1, 2, 3, 4, 5, 6],
    "column float64 A" => [23.654, 0.319, 10.0049, 89.01999, -3.41501, 52.0766],
    "column options"   => [Some(28), Some(300), None, Some(2), Some(-30), None],
    "column float64 B" => [23.6499, 0.399, 10.0061, 89.0105, -3.4331, 52.099999],
).unwrap();
df
Out[12]:
shape: (6, 4)
┌─────────────────┬──────────────────┬────────────────┬──────────────────┐
│ column integers ┆ column float64 A ┆ column options ┆ column float64 B │
│ ---             ┆ ---              ┆ ---            ┆ ---              │
│ i32             ┆ f64              ┆ i32            ┆ f64              │
╞═════════════════╪══════════════════╪════════════════╪══════════════════╡
│ 1               ┆ 23.654           ┆ 28             ┆ 23.6499          │
│ 2               ┆ 0.319            ┆ 300            ┆ 0.399            │
│ 3               ┆ 10.0049          ┆ null           ┆ 10.0061          │
│ 4               ┆ 89.01999         ┆ 2              ┆ 89.0105          │
│ 5               ┆ -3.41501         ┆ -30            ┆ -3.4331          │
│ 6               ┆ 52.0766          ┆ null           ┆ 52.099999        │
└─────────────────┴──────────────────┴────────────────┴──────────────────┘
In [13]:
let columns_with_float64 = vec![
    "column float64 A",
    "column float64 B",
];
columns_with_float64
Out[13]:
["column float64 A", "column float64 B"]
In [ ]:
lazyframe
.with_columns([
    cols(col1, col2, ..., colN)
   .apply(|series| 
       some_function(series), 
       GetOutput::from_type(DataType::Float64)
   )
]);

Comments