已获取PyReadonlyArray2，如何将指定形状的Numpy数组转为Rust的Vec<[f64;1024]>

核心思路 #

借助rust-numpy关联的ndarray视图能力，逐行处理二维数组，将每行转换为固定大小的[f64; 1024]数组，最终收集为目标Vec类型。

具体代码实现 #

use numpy::{PyReadonlyArray2, ndarray};
use pyo3::prelude::*;

#[pyfunction]
fn convert_numpy_arr(arr: PyReadonlyArray2<f64>) -> PyResult<Vec<[f64; 1024]>> {
    // 验证数组第二维度是否符合要求
    let shape = arr.shape();
    if shape[1] != 1024 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("数组第二维度必须为1024，当前为{}", shape[1]),
        ));
    }

    // 将PyReadonlyArray2转为ndarray只读视图
    let arr_view = arr.as_array();

    // 遍历每行并转换为固定大小数组
    arr_view.rows()
        .map(|row| {
            row.to_vec()
                .try_into()
                .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    "行数据转换为固定大小数组失败"
                ))
        })
        .collect()
}

关键细节说明 #

• 形状校验：提前检查数组维度，避免后续类型转换触发 panic。
• 视图转换：as_array()将Python侧数组转为Rust可直接操作的ndarray视图，简化行遍历逻辑。
• 安全类型转换：try_into()依赖编译时长度检查，结合前置维度校验，转换逻辑完全安全。
• 非连续数组兼容：即使输入numpy数组是非连续内存布局（如切片生成），rows()仍能正常遍历，to_vec()会自动复制数据到连续内存。

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