如何从R调用返回Vec<Vec<f64>>的Rust并行抽样函数？

因R代码运行过慢，用Rust实现了从M个不同正态分布中抽样N次的并行函数，返回Vec<Vec<f64>>。通过rextendr库从R调用时，编译报错：Robj: From<Vec<Vec<f64>>> trait未实现，同时不确定当前调用方式是否为推荐方案。

Rust代码 #

use rand_distr::{Normal, Distribution};
use rayon::prelude::*;

fn rust_rprednorm(n: i32, means: Vec<f64>, sds: Vec<f64>) -> Vec<Vec<f64>> {

    let mut preds = vec![vec![0.0; n as usize]; means.len()];

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

    preds
}

R调用代码 #

code <- r"(
    use rand_distr::{Normal, Distribution};
    use rayon::prelude::*;

    #[extendr]
    fn rust_rprednorm(n: i32, means: Vec<f64>, sds: Vec<f64>) -> Vec<Vec<f64>> {

        let mut preds = vec![vec![0.0; n as usize]; means.len()];

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

        preds
    }
)"

rust_source(code = code, dependencies = list(`rand` = "0.8.5", `rand_distr` ="0.4.3", `rayon` = "1.6.1"))

错误信息 #

Error in `invoke_cargo()`:
! Rust code could not be compiled successfully. Aborting.
✖ error[E0277]: the trait bound `Robj: From<Vec<Vec<f64>>>` is not satisfied
 --&gt; src\lib.rs:6:5
  |
6 |     #[extendr]
  |     ^^^^^^^^^^ the trait `From<Vec<Vec<f64>>>` is not implemented for `Robj`
  |
  = help: the following other types implement trait `From<T>`:
            <Robj as From<&'a [T]>>
            <Robj as From<&Altrep>>
            <Robj as From<&Primitive>>
            <Robj as From<&Robj>>
            <Robj as From<&Vec<T>>>
            <Robj as From<&extendr_api::Complexes>>
            <Robj as From<&extendr_api::Doubles>>
            <Robj as From<&extendr_api::Environment>>
          and 71 others
  = note: this error originates in the attribute macro `extendr` (in Nightly builds, run with -Z macro-backtrace for more info)
✖ error: aborting due to previous error
Traceback:

1. source("inla_predictive_distribution_utils.R")
2. withVisible(eval(ei, envir))
3. eval(ei, envir)
4. eval(ei, envir)
5. rust_source(code = code, dependencies = list(rand = "0.8.5", 
 .     rand_distr = "0.4.3", rayon = "1.6.1"))
6. invoke_cargo(toolchain = toolchain, specific_target = specific_target, 
 .     dir = dir, profile = profile, quiet = quiet, use_rtools = use_rtools)
7. check_cargo_output(compilation_result, message_buffer, tty_has_colors(), 
 .     quiet)
8. ui_throw("Rust code could not be compiled successfully. Aborting.", 
 .     error_messages, call = call, glue_open = "{<{", glue_close = "}>}")
9. withr::with_options(list(warning.length = message_limit_bytes), 
 .     rlang::abort(message, class = "rextendr_error", call = call))
10. force(code)
11. rlang::abort(message, class = "rextendr_error", call = call)
12. signal_abort(cnd, .file)

1. 适配rextendr的类型转换规则 #

rextendr的#[extendr]宏要求返回类型能直接转换为R的Robj，Vec<Vec<f64>>没有默认实现该转换，推荐以下两种处理方式：

方式一：改用rextendr矩阵类型返回 #

直接将结果构造为extendr_api::Matrix<f64>，这是最简洁的方案，符合R的数据结构习惯：

use rand_distr::{Normal, Distribution};
use rayon::prelude::*;
use extendr_api::prelude::*;

#[extendr]
fn rust_rprednorm(n: i32, means: Vec<f64>, sds: Vec<f64>) -> Matrix<f64> {
    let m = means.len();
    let n_usize = n as usize;
    
    // 预先分配扁平化数组，减少内存开销
    let mut preds = vec![0.0; m * n_usize];
    
    preds.par_chunks_mut(n_usize).enumerate().for_each(|(i, chunk)| {
        let mut rng = rand::thread_rng();
        // 每个分布只初始化一次，避免重复计算
        let normal = Normal::new(means[i], sds[i]).unwrap();
        chunk.iter_mut().for_each(|val| {
            *val = normal.sample(&mut rng);
        });
    });
    
    // 构造矩阵：行数为分布数量，列数为抽样次数
    Matrix::new(m as i32, n, preds)
}

方式二：手动实现类型转换（不推荐） #

如果必须返回Vec<Vec<f64>>，可以手动实现From<Vec<Vec<f64>>> for Robj trait，将嵌套向量转为R的列表结构：

use extendr_api::{Robj, Error};

impl From<Vec<Vec<f64>>> for Robj {
    fn from(vv: Vec<Vec<f64>>) -> Self {
        vv.into_iter()
            .map(|v| v.into())
            .collect::<Vec<Robj>>()
            .into()
    }
}

2. 优化Rust抽样逻辑 #

原代码中每次抽样都重新创建Normal分布对象，存在不必要的性能开销，建议在每个线程中仅初始化一次分布（已在方式一的代码中体现）。

3. 推荐的rextendr调用方式 #

直接在R中嵌入Rust代码适合快速测试，但函数逻辑复杂时，更推荐将Rust代码封装为独立Cargo库，便于维护和版本控制：

1. 创建Rust库项目：

cargo new --lib r_rprednorm

2. 修改Cargo.toml添加依赖：

[package]
name = "r_rprednorm"
version = "0.1.0"
edition = "2021"

[dependencies]
extendr-api = "0.6"
rand = "0.8.5"
rand_distr = "0.4.3"
rayon = "1.6.1"

3. 在src/lib.rs中编写代码并导出模块：

use rand_distr::{Normal, Distribution};
use rayon::prelude::*;
use extendr_api::prelude::*;

#[extendr]
fn rust_rprednorm(n: i32, means: Vec<f64>, sds: Vec<f64>) -> Matrix<f64> {
    // 实现代码同方式一
    let m = means.len();
    let n_usize = n as usize;
    let mut preds = vec![0.0; m * n_usize];
    
    preds.par_chunks_mut(n_usize).enumerate().for_each(|(i, chunk)| {
        let mut rng = rand::thread_rng();
        let normal = Normal::new(means[i], sds[i]).unwrap();
        chunk.iter_mut().for_each(|val| {
            *val = normal.sample(&mut rng);
        });
    });
    
    Matrix::new(m as i32, n, preds)
}

// 导出R模块
extendr_module! {
    mod r_rprednorm;
    fn rust_rprednorm;
}

4. 在R中安装并调用：

library(rextendr)
devtools::install("path/to/r_rprednorm")
library(r_rprednorm)

# 测试调用
means <- c(1.0, 2.0, 3.0)
sds <- c(0.5, 0.3, 0.8)
result <- rust_rprednorm(100, means, sds)

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