如何在Cython中将malloc分配的大数组作为Python对象返回或保存？

你遇到的编译报错原因是C语言原生的二级指针int **属于C层面的内存结构，Python解释器无法直接识别转换为Python对象，所以无法直接返回。另外你采用二级指针逐行分配内存的方案还存在两个问题：一是内存零散分布，访问时缓存命中率低，运行速度慢；二是手动分配的内存没有释放逻辑，会造成内存泄漏。

下面给出两种可行的解决方案：

方案1：直接返回Numpy数组（推荐） #

利用Cython对Numpy的原生支持，直接创建Numpy数组并写入数据，内存由Numpy自动管理，返回的数组可以直接被其他Python函数使用，性能也比手动分配二级指针更高。

1.1 修改后的my_script.pyx代码 #

import cython
import numpy as np
cimport numpy as np
from libc.stdlib cimport rand, srand, RAND_MAX
from libc.time cimport time

# 全局仅执行一次随机数种子初始化
srand(time(NULL))

def my_array_func(int a_param):
    # 创建10000*10000的int32类型Numpy数组，内存由Numpy托管
    cdef np.ndarray[np.int32_t, ndim=2] my_array = np.zeros((10000, 10000), dtype=np.int32)
    # 获取数组底层连续内存的指针
    cdef np.int32_t *arr_ptr = <np.int32_t *>my_array.data
    cdef int j, k
    # 直接写入连续内存，性能更高
    for j in range(10000):
        for k in range(10000):
            arr_ptr[j * 10000 + k] = <np.int32_t>(<float>rand()/RAND_MAX * a_param)
    # 直接返回Numpy数组，Python侧可直接使用
    return my_array

1.2 编译用的setup.py代码 #

需要在编译配置中加入Numpy的头文件路径：

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules = cythonize("my_script.pyx"),
    include_dirs = [np.get_include()]
)

1.3 Python侧调用示例 #

编译完成后直接导入使用即可：

import my_script
import numpy as np

# 直接获取生成好的Numpy数组
arr = my_script.my_array_func(100)
# 后续可直接用于Python侧的分析逻辑
print(arr.shape) # 输出 (10000, 10000)

方案2：直接写入本地文件后续读取 #

如果不需要直接返回数组，也可以在Cython中将生成的数据直接写入二进制文件，后续Python脚本直接读取即可。

2.1 新增写文件功能的Cython代码 #

在my_script.pyx中新增如下函数：

from libc.stdio cimport fopen, fwrite, fclose

def generate_and_save(int a_param, str file_path):
    cdef np.ndarray[np.int32_t, ndim=2] my_array = np.zeros((10000, 10000), dtype=np.int32)
    cdef np.int32_t *arr_ptr = <np.int32_t *>my_array.data
    cdef int j, k
    for j in range(10000):
        for k in range(10000):
            arr_ptr[j * 10000 + k] = <np.int32_t>(<float>rand()/RAND_MAX * a_param)
    # 写入二进制文件
    cdef FILE* fp = fopen(file_path.encode("utf-8"), "wb")
    fwrite(arr_ptr, sizeof(np.int32_t), 10000 * 10000, fp)
    fclose(fp)

2.2 Python侧读取文件示例 #

import numpy as np

# 直接读取二进制文件并reshape为二维数组
arr = np.fromfile("生成的文件路径", dtype=np.int32).reshape(10000, 10000)

注意事项 #

• 10000x10000的int32数组大小约为400MB，无论是返回还是存储都不会有性能压力
• 原代码中缺少time函数的导入，上述方案已经补充了对应导入逻辑，避免编译报错

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