如何将包含csr稀疏矩阵的Python Dataset对象序列化保存为单个文件

实现方案 #

直接使用NumPy原生的.npz归档格式即可，它天然支持在单个文件中存储多组二进制数据，性能和你原来分文件存储的效率几乎一致，不需要手动处理字节流拼接、偏移量计算等容易出错的逻辑。

核心思路 #

• 提取CSR稀疏矩阵的三个核心属性indptr、indices、data，以及矩阵形状shape参数
• 将meta_data序列化为JSON字符串后转成字节数组
• 将以上所有内容统一存入同一个npz文件

完整代码实现 #

import scipy.sparse as sparse
import json
import numpy as np
from pathlib import Path

class Dataset: 
    def __init__(self, X: sparse.csr_matrix, meta_data: dict):
        self.X = X
        self.meta_data = meta_data

    def save_to_single_file(self, dest_file: str) -> None: 
        # 存储稀疏矩阵的核心参数
        np.savez_compressed(
            dest_file,
            indptr=self.X.indptr,
            indices=self.X.indices,
            data=self.X.data,
            shape=self.X.shape,
            # meta_data转json字节存储
            meta_data=json.dumps(self.meta_data).encode("utf-8")
        )
    
    @classmethod
    def load_from_single_file(cls, src_file: str) -> "Dataset":
        content = np.load(src_file)
        # 重构稀疏矩阵
        X = sparse.csr_matrix(
            (content["data"], content["indices"], content["indptr"]),
            shape=content["shape"]
        )
        # 解析meta_data
        meta_data = json.loads(content["meta_data"].decode("utf-8"))
        return cls(X=X, meta_data=meta_data)

    def save_to_dir(self, dir_name: str) -> None:
       save_dir = Path(dir_name)
       # 注：原代码参数顺序有误，修正为路径在前，矩阵在后
       sparse.save_npz(save_dir / "X.npz", self.X)
       
       with open(save_dir / "meta_data.json", "w") as f:
           json.dump(self.meta_data, f)

其他可选方案说明 #

你提到的手动拼接字节流的方案是可行的，但需要额外在文件头部存储各段数据的长度偏移量，否则读取时无法区分不同数据段的边界，还要自行处理数组类型、大小端兼容等问题，开发成本更高，稳定性也不如直接使用成熟的npz格式，除非有特殊的自定义存储需求，否则不推荐使用。

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