English 中文(简体)
How to apply a function to multiple columns of a polars DataFrame in Rust
原标题:

I d like to apply a user-define function which takes a few inputs (corresponding some columns in a polars DataFrame) to some columns of a polars DataFrame in Rust. The pattern that I m using is as below. I wonder is this the best practice?

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
}
问题回答

You can downcast the Series to the proper type you want to iterate over, and then use rust iterators to apply your logic.

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"),
    }
}

Lazy API

You don t have to leave the lazy API to be able to access my_black_box_function.

We can collect the columns we want to apply in a Struct data type and then apply a closure over that Series.

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()
}

The solution I found working for me is with map_multiple(my understanding - this to be used if no groupby/agg) or apply_multiple(my understanding - whenerver you have groupby/agg). Alternatively, you could also use map_many or apply_many. See below.

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

fn main() {
    let df = df! [
        "names" => ["a", "b", "a"],
        "values" => [1, 2, 3],
        "values_nulls" => [Some(1), None, Some(3)],
        "new_vals" => [Some(1.0), None, Some(3.0)]
    ].unwrap();

    println!("{:?}", df);

    //df.try_apply("values_nulls", |s: &Series| s.cast(&DataType::Float64)).unwrap();

    let df = df.lazy()
        .groupby([col("names")])
        .agg( [
            total_delta_sens().sum()
        ]
        );

    println!("{:?}", df.collect());
}

pub fn total_delta_sens () -> Expr {
    let s: &mut [Expr] = &mut [col("values"), col("values_nulls"),  col("new_vals")];

    fn sum_fa(s: &mut [Series])->Result<Series>{
        let mut ss = s[0].cast(&DataType::Float64).unwrap().fill_null(FillNullStrategy::Zero).unwrap().clone();
        for i in 1..s.len(){
            ss = ss.add_to(&s[i].cast(&DataType::Float64).unwrap().fill_null(FillNullStrategy::Zero).unwrap()).unwrap();
        }
        Ok(ss) 
    }

    let o = GetOutput::from_type(DataType::Float64);
    map_multiple(sum_fa, s, o)
}

Here total_delta_sens is just a wrapper function for convenience. You don t have to use it.You can do directly this within your .agg([]) or .with_columns([]) : lit::<f64>(0.0).map_many(sum_fa, &[col("norm"), col("uniform")], o)

Inside sum_fa you can as Richie already mentioned downcast to ChunkedArray and .iter() or even .par_iter() Hope that helps

With Polars version = "0.30", use:

.with_columns([
    cols( col1, col2, ..., colN)
   .apply(|series| 
       some_function(), 
       GetOutput::from_type(DataType::Float64)
   )
]);

The Cargo.toml:

[dependencies]
polars = { version = "0.30", features = [
    "lazy", # Lazy API
    "round_series", # round underlying float types of Series
] }

And the main() function:

use std::error::Error;

use polars::{
    prelude::*,
    datatypes::DataType,
};

fn main()-> Result<(), Box<dyn Error>> {

    let dataframe01: DataFrame = 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],
    )?;

    println!("dataframe01: {dataframe01}
");

    let columns_with_float64: Vec<&str> = vec![
        "column float64 A",
        "column float64 B",
    ];

    // Format only the columns with float64

    let lazyframe: LazyFrame = dataframe01
        .lazy()
        .with_columns([
            cols(columns_with_float64)
            .apply(|series| 
                Ok(Some(series.round(2)?)), 
                GetOutput::from_type(DataType::Float64)
            )
         ]);
    
    let dataframe02: DataFrame = lazyframe.collect()?;
    
    println!("dataframe02: {dataframe02}
");

    let series_a: Series = Series::new("column float64 A", &[23.65, 0.32, 10.00, 89.02, -3.42, 52.08]);
    let series_b: Series = Series::new("column float64 B", &[23.65,  0.4, 10.01, 89.01, -3.43, 52.1]);

    assert_eq!(dataframe02.column("column float64 A")?, &series_a);
    assert_eq!(dataframe02.column("column float64 B")?, &series_b);

    Ok(())
}

The output:

dataframe01: 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        │
└─────────────────┴──────────────────┴────────────────┴──────────────────┘

dataframe02: shape: (6, 4)
┌─────────────────┬──────────────────┬────────────────┬──────────────────┐
│ column integers ┆ column float64 A ┆ column options ┆ column float64 B │
│ ---             ┆ ---              ┆ ---            ┆ ---              │
│ i32             ┆ f64              ┆ i32            ┆ f64              │
╞═════════════════╪══════════════════╪════════════════╪══════════════════╡
│ 1               ┆ 23.65            ┆ 28             ┆ 23.65            │
│ 2               ┆ 0.32             ┆ 300            ┆ 0.4              │
│ 3               ┆ 10.0             ┆ null           ┆ 10.01            │
│ 4               ┆ 89.02            ┆ 2              ┆ 89.01            │
│ 5               ┆ -3.42            ┆ -30            ┆ -3.43            │
│ 6               ┆ 52.08            ┆ null           ┆ 52.1             │
└─────────────────┴──────────────────┴────────────────┴──────────────────┘





相关问题
Row/column counter in apply functions

What if one wants to apply a functon i.e. to each row of a matrix, but also wants to use as an argument for this function the number of that row. As an example, suppose you wanted to get the n-th root ...

Is R s apply family more than syntactic sugar?

...regarding execution time and / or memory. If this is not true, prove it with a code snippet. Note that speedup by vectorization does not count. The speedup must come from apply (tapply, sapply, ......

PickUp/Apply Paragraph Formatting in PowerPoint 2007

In PowerPoint 2007, PickUp/Apply does not capture some paragraph formatting, such as bullet formatting, when used programmatically (VBA). Adding the PickUp and Apply buttons to the Quick Access ...

How to make a list of arrays, not their symbols, in Lisp?

I m trying to make a function to get a delta between arrays, but right now just want to make a subset: get Nth element. (defvar p1 #(1 2)) (defvar p2 #(3 4)) (mapcar (lambda (x) (aref x 0)) (p1 ...

jQuery: inherit functions to several objects

i made several table-based-widgets (listview-kind-of) which all have the same characteristics: styling odd/even rows, hover on/off, set color onClick, deleting a row when clicking on trash-icon. so ...

热门标签