array库与numpy选型：大量float向量磁盘存储方案优化问询

一、先明确：JSON方案完全不适合该场景 #

JSON是文本格式，存储float时会把二进制数值转成字符串，磁盘占用比二进制格式大3-5倍，读写速度慢，完全不适合大量数值型数据的存储，直接排除。

二、array库 vs Numpy：速度与内存对比 #

你观察到两者生成的二进制文件一致，是因为它们都采用了标准的IEEE 754双精度浮点数格式（'d'对应float64），但底层实现的开销有差异：

• 速度：array.array的写入速度通常略快于Numpy。因为Numpy数组包含额外的元数据（形状、维度、数据类型信息等），在转换和写入时会有少量额外开销；而array.array是更轻量的连续内存结构，更接近C语言数组，初始化和写入的开销更低。
• 内存占用：array.array的内存占用略低于Numpy数组，同样是因为Numpy的数组对象携带更多元数据。不过如果你的输入是list[list[float]]，两者在转换时都需要将数据复制到连续内存块，这一步的内存开销基本相当。

三、核心优化方案 #

1. 避免临时内存开销：用生成器展平数据 #

如果你的输入是嵌套的list[list[float]]，直接传入array.array或Numpy构造函数会先创建一个巨大的一维临时list，浪费内存。改用生成器逐个yield元素，避免临时list的内存占用：

array库优化版

import array

def save_vectors_array(vectors, file_name):
    if not vectors:
        return
    vec_length = len(vectors[0])
    total_count = len(vectors)
    
    # 用生成器展平所有向量，避免创建大临时list
    flat_generator = (item for vec in vectors for item in vec)
    
    with open(file_name, 'wb') as f:
        # 先写入元数据：单个向量长度、总向量数（方便读取时还原结构）
        meta_arr = array.array('I', [vec_length, total_count])
        meta_arr.tofile(f)
        
        # 写入核心数据
        data_arr = array.array('d', flat_generator)
        data_arr.tofile(f)

Numpy优化版

import numpy as np

def save_vectors_numpy(vectors, file_name):
    if not vectors:
        return
    vec_length = len(vectors[0])
    total_count = len(vectors)
    total_elements = total_count * vec_length
    
    # 用fromiter直接从生成器构造数组，避免临时list
    flat_array = np.fromiter(
        (item for vec in vectors for item in vec),
        dtype=np.float64,
        count=total_elements
    )
    # 重塑为二维数组（视图操作，无额外内存开销）
    vec_array = flat_array.reshape(total_count, vec_length)
    
    with open(file_name, 'wb') as f:
        # 写入元数据
        np.array([vec_length, total_count], dtype=np.uint32).tofile(f)
        # 写入向量数据
        vec_array.tofile(f)

2. 超大数据量：用内存映射或分块写入 #

如果数据量远超内存容量（比如TB级），可以用Numpy的内存映射文件，直接在磁盘上操作数组，无需加载全部数据到内存：

import numpy as np

def save_vectors_mmap(vectors, file_name):
    if not vectors:
        return
    vec_length = len(vectors[0])
    total_count = len(vectors)
    
    # 创建内存映射数组，仅占用分块大小的内存
    mmap_arr = np.memmap(
        file_name,
        dtype=np.float64,
        mode='w+',
        shape=(total_count, vec_length)
    )
    
    # 分块写入，控制单批次内存占用
    chunk_size = 10000  # 根据内存情况调整
    for i in range(0, total_count, chunk_size):
        end_idx = min(i + chunk_size, total_count)
        mmap_arr[i:end_idx] = vectors[i:end_idx]
    
    # 删除引用，确保数据刷入磁盘
    del mmap_arr

3. 兼顾压缩与随机访问：用HDF5格式 #

如果需要压缩磁盘空间，同时支持随机读取任意向量，推荐使用h5py库操作HDF5格式：

import h5py

def save_vectors_h5py(vectors, file_name):
    if not vectors:
        return
    vec_length = len(vectors[0])
    total_count = len(vectors)
    
    with h5py.File(file_name, 'w') as f:
        # 创建数据集，支持压缩（可选，比如gzip）
        dset = f.create_dataset(
            'vectors',
            shape=(total_count, vec_length),
            dtype='f8',
            compression='gzip'  # 开启压缩，节省磁盘空间
        )
        
        # 分块写入（内存不足时使用）
        chunk_size = 10000
        for i in range(0, total_count, chunk_size):
            end_idx = min(i + chunk_size, total_count)
            dset[i:end_idx] = vectors[i:end_idx]

HDF5的优势是支持超大规模数据、压缩、随机访问，且后续可以直接用Numpy或Pandas读取处理。

四、最终方案选择 #

• 极致速度+低内存：优先选择array.array配合生成器的方案，开销最小，读写最快。
• 后续需Numpy处理：用Numpy的fromiter优化版，避免额外的数组转换开销。
• 超大数据量：用Numpy内存映射或HDF5分块写入。
• 需要压缩/随机访问：选择HDF5格式。

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