r - Polars Rust Melt() 明显慢于 R stack()

Question

我有一些 R 代码，它需要一个宽的 data.frame 并将其堆叠到一个窄的 data.frame 中。我用 Rust 重写了这个，但发现它非常慢。我想知道我是否使用了不好的做法或这里的某些东西会扼杀速度。

原始 R 版本:

df = cbind(df[ncol(df)], df[ncol(df)-3], df[ncol(df)-2], df[ncol(df)-1], stack(df[1:(ncol(df)-4)]))

stack(df[1:(ncol(df)-4)])部分获取除最后 4 列(通常为 1,000 列)之外的所有列并将它们堆叠起来。它还创建第二列，指示行来自哪一列。然后我将其他 4 列重新绑定(bind)到它。 R 自动重复它们以匹配窄 df 的新长度。

这是我的 Polars eager 版本:

let n = 1000;
let sample_cols = (0..n).collect::<Vec<i32>>()
    .par_iter()
    .map(|l| format!("{}", l))
    .collect::<Vec<String>>();

let mut df = df.melt(&["A", "B", "C", "D"], sample_cols).unwrap();

sample_cols 是一个 Vec，包含要堆叠的列名称，这些列名称是 0 到 999 之间的字符串，适用于 1000 个样本。

这是惰性版本:

let n = 1000;
let sample_cols = (0..n).collect::<Vec<i32>>()
    .par_iter()
    .map(|l| format!("{}", l))
    .collect::<Vec<String>>();

let melt_args = MeltArgs {
    id_vars: vec!["A".into(), "B".into(), "C".into(), "D".into()],
    value_vars: sample_cols,
    variable_name: None,
    value_name: None,
};

let mut df = df.lazy().melt(melt_args).collect()?;

两个 Rust 版本的速度相似，但比 R 慢得多。当 n = 100,000 时，R 代码平均需要 0.45 秒，但有时只需 0.23 秒，而两个 Rust 版本需要 13.5 秒到 14.5 秒。

如果您想自己运行它，这应该生成虚拟数据并运行它，只需确保一次仅使用 eager 或 lazy 版本:

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

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
}

let n = 100000;

let means: Vec<f64> = vec![0.0; 15];
let sds: Vec<f64> = vec![1.0; 15];
let preds = rprednorm(n as usize, means, sds);

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>>()
    )?;

let start = Instant::now();
let sample_cols= (0..n).collect::<Vec<i32>>()
    .par_iter()
    .map(|l| format!("{}", l))
    .collect::<Vec<String>>();

df.with_column(Series::new("A", &["1", "2", "3", "1", "2", "3'", "1", "2", "3", "1", "2", "3", "1", "2", "3"]));
df.with_column(Series::new("B", &["1", "1", "1", "2", "2", "2", "3", "3", "3", "4", "4", "4", "5", "5", "5"]));
df.with_column(Series::new("C", &["1", "2", "3", "1", "2", "3'", "1", "2", "2", "1", "2", "3'", "1", "2", "3"]));
df.with_column(Series::new("D", (0..df.shape().0 as i32).collect::<Vec<i32>>()));

let melt_args = MeltArgs {
    id_vars: vec!["A".into(), "B".into(), "C".into(), "D".into()],
    value_vars: sample_cols,
    variable_name: None,
    value_name: None,
};

let start = Instant::now();
let mut df = df.lazy().melt(melt_args).collect()?;
let duration = start.elapsed();
println!("{:?}", duration);

let start = Instant::now();
let mut df = df.melt(&["A", "B", "C", "D"], &sample_cols).unwrap();
let duration = start.elapsed();
println!("{:?}", duration);

Answer 1

我在Github上提交了一个问题，现有的实现从O(n^2)改进到O(n)，现在比R更快。它不是最新更新的一部分，所以你需要安装来自 github 而不是 crates.io