Polars Rust版melt()性能远低于R语言stack()问题排查

我有一段R代码可快速将宽格式data.frame转换为窄格式，但用Rust的Polars库重写后运行速度极慢，怀疑是用法不当导致的性能问题。

原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列外的所有列（通常为千级甚至更多）并堆叠，同时生成标记来源列的新列；R会自动重复最后4列以匹配窄表长度。

Rust Polars 实现版本 #

立即执行版本 #

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

延迟执行版本 #

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

说明：sample_cols是包含待堆叠列名的Vec，内容为0到n-1的字符串。

性能对比 #

当n=100000时，两者性能差距显著：

• R代码：平均耗时0.45s，最快0.23s
• Rust Polars版本（立即/延迟执行速度相近）：耗时13.5-14.5s

测试数据生成代码 #

以下是可复现问题的测试数据生成代码（每次仅使用一种执行版本）：

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
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let n = 100000;

    let means: Vec<f64> = vec![0.0; 15];
    let sds: Vec<f64> = vec![1.0; 15];
    let preds = multi_rnorm(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>>()
    )?;

    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 sample_cols= (0..n).collect::<Vec<i32>>()
        .par_iter()
        .map(|l| format!("{}", l))
        .collect::<Vec<String>>();

    // 测试延迟执行版本
    let start = Instant::now();
    let melt_args = MeltArgs {
        id_vars: vec!["A".into(), "B".into(), "C".into(), "D".into()],
        value_vars: sample_cols.clone(),
        variable_name: None,
        value_name: None,
    };
    let _df_lazy = df.lazy().melt(melt_args).collect()?;
    let duration = start.elapsed();
    println!("延迟执行耗时: {:?}", duration);

    // 测试立即执行版本
    let start = Instant::now();
    let _df_eager = df.melt(&["A", "B", "C", "D"], &sample_cols).unwrap();
    let duration = start.elapsed();
    println!("立即执行耗时: {:?}", duration);

    Ok(())
}

说明：修正了原代码中的函数名错误（rprednorm改为multi_rnorm），补充了main函数的错误处理逻辑，确保代码可编译运行。

性能优化建议 #

1. 避免不必要的字符串列名生成
   原代码中通过par_iter()批量生成字符串列名，带来额外内存开销和计算时间。可直接通过列名索引范围选择待堆叠列：

   let all_cols = df.get_column_names();
   let value_vars = &all_cols[0..all_cols.len()-4];
   

2. 使用Polars原生随机数据生成
   原测试代码用ndarray生成数据再转Polars DataFrame，存在不必要的内存拷贝。可直接用Polars原生API生成数据：

   let mut df = DataFrame::new(Vec::new())?;
   for i in 0..n {
       let s = Series::random(&format!("{}", i), 15, Distribution::Normal(0.0, 1.0))?;
       df.with_column(s)?;
   }
   

3. 优化ID列数据类型
   原代码中ID列（A、B、C）使用字符串类型，对于重复值较多的列，可转换为Categorical类型减少内存占用，提升melt操作效率：

   df.with_column(Series::new("A", &["1", "2", "3", ...]).cast(&DataType::Categorical(None))?)?;
   

4. 升级Polars版本
   确保使用最新版本的Polars，新版本通常会修复性能问题并优化melt操作的底层实现。

内容的提问来源于stack exchange，提问作者GBPU