Quickly create polars data frame from 2D ndarray

Question

I am sampling from m normal distributions n times each, and storing the samples in a 2-dimensional ndarray. I then need to turn this into a polars dataframe as I need to add some string columns and do a lot of querying/aggregating.

I think where I convert the array to a dataframe is inefficient. The Polars docs did not seem to have any examples of doing this, every example I saw uses small amounts of data typed out by hand, whereas I have 10,000 samples each from 100,000 normal distributions.

Here is the function that samples:

use rand_distr::{Normal, Distribution};
use polars::prelude::*;
use rayon::prelude::*;
use ndarray::Array2;
use cute::c;
#[macro_use]
extern crate fstrings;

fn multi_rnorm(n: usize, means: Vec<f64>, sds: Vec<f64>) -> Array2<f64> {

    let mut preds: Array2<f64> = Array2::zeros((means.len(), n));

    preds.axis_iter_mut(ndarray::Axis(0)).into_par_iter().enumerate().for_each(|(i, mut row)| {

        let mut rng = rand::thread_rng();
        (0..n).into_iter().for_each(|j| {
            let normal = Normal::new(means[i], sds[i]).unwrap();
            row[j as usize] = normal.sample(&mut rng);
        })
    });
    preds
}

And here is how I am creating the data frame:

let mut df: DataFrame = DataFrame::new(
    c![Series::new(
                &f!("{tuple.0}"), 
                tuple.1.to_vec()), for tuple in preds.axis_iter(ndarray::Axis(1))
                    .into_iter()
                    .enumerate()
                    .collect::<Vec<_>>()])
                    .unwrap();

Main looks like this:

fn main() {

    let means = vec![0.0; 99_128];
    let sds = vec![1.0; 99_128];

    let preds = rprednorm(10000, means, sds);

    let mut df: DataFrame = DataFrame::new(
        c![Series::new(
                    &f!("{tuple.0}"), 
                    tuple.1.to_vec()), for tuple in preds.axis_iter(ndarray::Axis(1))
                        .into_iter()
                        .enumerate()
                        .collect::<Vec<_>>()])
                        .unwrap();
}

I am using the c! macro from the cute crate to emulate python dictionary comprehensions, as well as the f!() macro from fstrings to perform string interpolation.

I have tried substituting into_par_iter() for into_iter(), but there is no speedup.

I was very surprised to not see any examples of how to do this in Polars, I figured there would be some implementation of FromNdarray so you could just call DataFrame::new(the_ndarray) or something.

Does my implementation look any good, or can it be faster?

Answer 1

It is possible to optimize the creation of the DataFrame from the Array2 by 5x by parallelizing it.

let mut df: DataFrame = DataFrame::new(
    preds.axis_iter(ndarray::Axis(1))
        .into_par_iter()
        .enumerate()
        .map(|(i, col)| {
            Series::new(
                &f!("{i}"),
                col.to_vec()
            )
        })
        .collect::<Vec<Series>>()
    ).unwrap();