Using for loop on a polars DataFrame

code review
1

I want to use polars DataFrames to store data and create new columns via looping over the rows in the DataFrame (I am doing it in Rust so it is faster than Python). The following code works fine, but I want to get opinions on if it can be improved any. For example, it seems to need a lot of unwrap() calls. Also, I would like to the code to be as fast and efficient as possible.

As the example below shows, the signal column gets updated based on some conditions that are checked in the for loop. Then the df is updated with the signal column. I am mainly using this code for metrics that don't lend themselves to vectorization. Or should I choose a different approach because polars is not well suited for row-wise looping?

use polars::prelude::*;

fn main() {

    let col1 = Series::new("col1", vec![1, 2, 3, 4, 5]);
    let col2 = Series::new("col2", vec![5, 4, 3, 2, 1]);
    let col3 = Series::new("col3", vec![1, 2, 3, 4, 5]);
    let mut df = DataFrame::new(vec![col1, col2, col3]).unwrap();

    let mut signal = vec![0; df.height()];
    
    // Iterate over DataFrame and update signal
    for i in 0..df.height() {
        if df.column("col1").unwrap().get(i).unwrap() == AnyValue::Int32(3) 
            && df.column("col2").unwrap().get(i).unwrap() == AnyValue::Int32(3)
            && df.column("col3").unwrap().get(i).unwrap() == AnyValue::Int32(3) 
        {
            signal[i] = 1;
        }
    }
    
    // Add new column to DataFrame
    let new_col = Series::new("signal", signal.clone());
    
    let df = df.with_column(new_col).unwrap();
    println!("after: {:?}", df);
}
2 
    let expr = (col("col1").eq(3))
        .and(col("col2").eq(3))
        .and(col("col3").eq(3));
    println!(
        "{:?}",
        df.clone()
            .lazy()
            .select([
                col("*"),
                expr.clone().alias("col1 == col2 == col3 == 3"),
                when(expr).then(1).otherwise(0).alias("signal")
            ])
            .collect()
            .unwrap()
    );

rust explorer

2 Likes
3

That is a big improvement from my version. Thank you so much.