使用PyO3从Rust调用Python函数的最优返回类型

1. Array2 vs Vec<Vec>的性能差异 #

Array2（来自ndarray crate）的性能远优于Vec<Vec<usize>>。原因在于：

• Vec<Vec<>>是嵌套堆分配，每行内存不连续，遍历时会频繁触发缓存失效，拖慢速度；
• Array2采用连续内存布局，默认C顺序下整行数据在内存中是连续的，缓存命中率更高，大数量级遍历的速度提升非常明显。

2. 从numpy.ndarray转换为Array2的实现 #

首先在Cargo.toml中添加依赖：

[dependencies]
pyo3 = { version = "0.21", features = ["numpy"] }
ndarray = "0.15"
pyo3-numpy = "0.21"

然后在Rust代码中完成转换：

use pyo3::prelude::*;
use pyo3::types::PyArray2;
use ndarray::Array2;

#[pyfunction]
fn rust_func(py: Python, my_py_func: &PyAny) -> PyResult<()> {
    // 调用Python函数并提取numpy数组（注意类型匹配原dtype为int64，对应Rust的i64）
    let py_array: &PyArray2<i64> = my_py_func.call0()?.extract()?;
    
    // 零拷贝获取数组视图（性能最优，只读场景推荐）
    let array_view = py_array.as_array();
    // 如果需要修改数组，再拷贝为拥有所有权的Array2
    let mut array2 = array_view.to_owned();
    
    // 遍历行操作示例
    for row in array2.rows() {
        let val1 = row[0];
        let val2 = row[1];
        // 你的业务逻辑
    }

    Ok(())
}

注意：若原numpy数组是Fortran顺序（列优先），可通过.as_standard_layout()转为C顺序，进一步优化行遍历的缓存效率。

3. 适合output_array的高性能数据结构选择 #

根据后续操作场景选择：

• 只读遍历行：优先用ArrayView2<i64>（零拷贝，直接复用Python数组内存），其次是Array2<i64>（连续内存，拷贝一次后遍历效率拉满）；
• 需要修改数组：Array2<i64>是最优解，连续内存布局在修改时同样能保证缓存效率；
• 需与其他数值库交互：如果要对接nalgebra等线性代数库，可转换为nalgebra::Matrix<i64, Dynamic, Const<2>>，但ndarray与PyO3的生态适配性更好；
• 类型注意事项：原numpy数组是int64，对应Rust的i64，不要随意用usize，避免跨架构的类型不匹配问题，除非你能确保数值范围不会超出usize的限制。

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