如何高效存储大量CSR稀疏数组？解决HDF5+Pickle报错问题

最初尝试将CSR矩阵转为密集数组后存入HDF5，但因一次性分配内存导致内存不足：

from scipy.sparse import csr_matrix
import h5py

As = [ csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]]),
       csr_matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
       csr_matrix([[2, 0, 0], [0, 3, 0], [0, 0, 4]]) ]

np_matrices = [mat.toarray() for mat in As]

with h5py.File(filename, "w") as f:
    f.create_dataset("matrices", data=np_matrices)

尝试用Pickle序列化后存入HDF5时，出现错误：VLEN strings do not support embedded NULLs，代码如下：

import pickle
import h5py

pickled_obj = pickle.dumps(As)

with h5py.File('obj.hdf5', 'w') as f:
    dset = f.create_dataset('obj', data=pickled_obj) 

求解决该错误或更高效的存储方案。

方法1：拆分CSR核心组件存入HDF5（推荐） #

CSR矩阵的核心是data、indices、indptr三个数组，外加shape属性。可以把所有矩阵的这些组件分别存入HDF5，同时记录每个矩阵的shape，既节省内存又避免序列化问题。

示例代码：

from scipy.sparse import csr_matrix, csr_array
import h5py
import numpy as np

As = [ csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]]),
       csr_matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
       csr_matrix([[2, 0, 0], [0, 3, 0], [0, 0, 4]]) ]

# 收集所有矩阵的核心数据
all_data = []
all_indices = []
all_indptr = []
all_shapes = []
data_offset = 0

for mat in As:
    all_data.extend(mat.data)
    all_indices.extend(mat.indices)
    # 给当前矩阵的indptr加上全局数据偏移，避免索引冲突
    all_indptr.extend(mat.indptr + data_offset)
    all_shapes.append(mat.shape)
    data_offset += len(mat.data)

# 转为numpy数组
all_data = np.array(all_data)
all_indices = np.array(all_indices)
all_indptr = np.array(all_indptr)
all_shapes = np.array(all_shapes)

# 写入HDF5
with h5py.File('csr_matrices.h5', 'w') as f:
    f.create_dataset('data', data=all_data)
    f.create_dataset('indices', data=all_indices)
    f.create_dataset('indptr', data=all_indptr)
    f.create_dataset('shapes', data=all_shapes)

# 读取还原
with h5py.File('csr_matrices.h5', 'r') as f:
    data = f['data'][:]
    indices = f['indices'][:]
    indptr = f['indptr'][:]
    shapes = f['shapes'][:]

restored_As = []
ptr_pos = 0
for shape in shapes:
    # 提取当前矩阵的indptr并还原偏移
    mat_indptr = indptr[ptr_pos:ptr_pos + shape[0]+1] - indptr[ptr_pos]
    # 提取对应的数据和索引
    mat_data = data[indptr[ptr_pos]:indptr[ptr_pos + shape[0]]]
    mat_indices = indices[indptr[ptr_pos]:indptr[ptr_pos + shape[0]]]
    # 重建CSR矩阵
    restored_mat = csr_array((mat_data, mat_indices, mat_indptr), shape=shape)
    restored_As.append(restored_mat)
    ptr_pos += shape[0] + 1

方法2：修正Pickle+HDF5的存储方式 #

之前的错误是因为HDF5默认将Pickle二进制数据识别为字符串，而Pickle包含NULL字节。只需创建字节型数据集即可解决：

示例代码：

import pickle
import h5py
import numpy as np

As = [ csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]]),
       csr_matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
       csr_matrix([[2, 0, 0], [0, 3, 0], [0, 0, 4]]) ]

pickled_obj = pickle.dumps(As)

# 写入HDF5（指定字节类型）
with h5py.File('obj.hdf5', 'w') as f:
    dset = f.create_dataset('obj', shape=(1,), dtype=h5py.special_dtype(vlen=bytes))
    dset[0] = pickled_obj

# 读取还原
with h5py.File('obj.hdf5', 'r') as f:
    pickled_data = f['obj'][0]
    restored_As = pickle.loads(pickled_data)

方法3：用Zarr分块存储稀疏矩阵 #

Zarr支持分块存储，适合大规模数据，且原生支持稀疏数组（需安装zarr和sparse库）：

示例代码：

from scipy.sparse import csr_matrix
import zarr
from sparse import COO

As = [ csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]]),
       csr_matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
       csr_matrix([[2, 0, 0], [0, 3, 0], [0, 0, 4]]) ]

# 创建Zarr存储目录
store = zarr.DirectoryStore('csr_zarr_store')
root = zarr.group(store=store)

# 逐个存入矩阵（CSR转COO格式）
for i, mat in enumerate(As):
    coo_mat = COO.from_scipy_sparse(mat)
    root.create_dataset(f'matrix_{i}', data=coo_mat)

# 读取还原
restored_As = []
for i in range(len(As)):
    coo_mat = root[f'matrix_{i}'][:]
    restored_mat = coo_mat.to_scipy_sparse('csr')
    restored_As.append(restored_mat)

方法4：批量写入单个NPZ文件 #

手动将多个矩阵的组件打包成字典，存入NPZ文件，轻量且高效：

示例代码：

from scipy.sparse import csr_matrix
import numpy as np

As = [ csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]]),
       csr_matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
       csr_matrix([[2, 0, 0], [0, 3, 0], [0, 0, 4]]) ]

# 构建存储字典
save_dict = {}
for i, mat in enumerate(As):
    save_dict[f'data_{i}'] = mat.data
    save_dict[f'indices_{i}'] = mat.indices
    save_dict[f'indptr_{i}'] = mat.indptr
    save_dict[f'shape_{i}'] = np.array(mat.shape)

# 写入NPZ
np.savez('multiple_csr.npz', **save_dict)

# 读取还原
loaded_dict = np.load('multiple_csr.npz')
restored_As = []
num_matrices = len([k for k in loaded_dict.keys() if k.startswith('data_')])

for i in range(num_matrices):
    data = loaded_dict[f'data_{i}']
    indices = loaded_dict[f'indices_{i}']
    indptr = loaded_dict[f'indptr_{i}']
    shape = tuple(loaded_dict[f'shape_{i}'])
    restored_mat = csr_matrix((data, indices, indptr), shape=shape)
    restored_As.append(restored_mat)

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