如何在磁盘存储海量numpy矩阵并支持多线程写入？

针对你的需求——存储数百万个shape=(224,224,1)的np.float32矩阵，同时需要高效访问、透明压缩和多线程写入支持（自行处理冲突），我整理了几个在工业界和科研中常用的可靠方案：

1. TileDB：面向数组的分布式存储（最匹配数组场景） #

TileDB是专门为多维数组设计的存储系统，完美契合numpy矩阵的存储需求，核心优势包括：

• 透明压缩：原生支持Blosc、Gzip、Zstd等多种压缩算法，配置后无需手动处理压缩/解压，完全透明。
• 并行写入支持：原生支持多线程/多进程写入，你可以通过分片写入（给每个线程分配独立的数组切片或索引范围）来避免冲突，或者利用其事务机制自定义冲突处理逻辑（比如覆盖、跳过冲突数据）。
• 高效访问：支持快速的随机访问和批量读取，性能媲美HDF5，同时比HDF5更适合并发场景。

简单使用示例：

import tiledb
import numpy as np
from concurrent.futures import ThreadPoolExecutor

# 定义数组 schema：shape 设为总矩阵数 × 224 × 224 × 1
schema = tiledb.ArraySchema(
    domain=tiledb.Domain(
        tiledb.Dim(name="idx", domain=(0, 999999), tile=1000, dtype=np.int64),
        tiledb.Dim(name="h", domain=(0, 223), tile=224, dtype=np.int32),
        tiledb.Dim(name="w", domain=(0, 223), tile=224, dtype=np.int32),
        tiledb.Dim(name="c", domain=(0, 0), tile=1, dtype=np.int32),
    ),
    attrs=[tiledb.Attr(name="data", dtype=np.float32, compression="blosc")],
    sparse=False
)

# 创建数组存储
tiledb.Array.create("matrix_store.tiledb", schema)

# 多线程写入：每个线程负责一个独立的批次范围（自行确保无重叠，避免冲突）
def write_batch(start_idx, end_idx):
    with tiledb.open("matrix_store.tiledb", "w") as arr:
        batch_data = np.random.rand(end_idx - start_idx, 224, 224, 1).astype(np.float32)
        arr[start_idx:end_idx] = batch_data

with ThreadPoolExecutor(max_workers=4) as executor:
    executor.submit(write_batch, 0, 250000)
    executor.submit(write_batch, 250000, 500000)
    executor.submit(write_batch, 500000, 750000)
    executor.submit(write_batch, 750000, 1000000)

2. Zarr + Dask：轻量级分块数组存储（灵活易扩展） #

Zarr是一种基于chunk的数组存储格式，和HDF5类似但更适合并发场景，搭配Dask可以轻松实现并行写入：

• 透明压缩：默认支持Blosc压缩，也可配置其他算法，压缩和解压完全透明。
• 多线程/进程写入：Zarr允许多个线程同时写入不同的chunk（只要你确保线程之间的chunk无重叠，自行处理冲突），Dask可以帮你自动管理并行任务，也可以手动控制线程的写入范围。
• 高效访问：支持随机访问任意chunk，读取速度快，且可以直接和numpy数组互转。

简单使用示例：

import zarr
import numpy as np
from concurrent.futures import ThreadPoolExecutor

# 创建Zarr数组，设置chunk大小（比如每个chunk存1000个矩阵）
store = zarr.DirectoryStore("matrix_zarr")
zarr_array = zarr.create(
    shape=(1000000, 224, 224, 1),
    dtype=np.float32,
    chunks=(1000, 224, 224, 1),
    compression="blosc",
    store=store
)

# 多线程写入：每个线程负责一个chunk的范围
def write_chunk(start_idx):
    end_idx = start_idx + 1000
    zarr_array[start_idx:end_idx] = np.random.rand(1000, 224, 224, 1).astype(np.float32)

with ThreadPoolExecutor(max_workers=4) as executor:
    for idx in range(0, 1000000, 1000):
        executor.submit(write_chunk, idx)

3. LMDB：内存映射键值存储（极致读写性能） #

如果你倾向于用键值对的方式存储每个矩阵（比如用唯一ID作为键），LMDB是绝佳选择：

• 极致性能：基于内存映射的键值存储，读写速度极快，适合数百万个小数据的场景，避免了大量小文件的inode问题。
• 多线程写入支持：LMDB支持多线程事务，你可以通过独立事务或分区键前缀来处理冲突（比如给每个线程分配不同的键前缀，避免键冲突）。
• 透明压缩：LMDB本身不内置压缩，但你可以封装一层逻辑，写入前用Blosc或numpy的savez_compressed压缩数据，读取时自动解压，实现透明压缩效果。

简单使用示例：

import lmdb
import numpy as np
import blosc
from concurrent.futures import ThreadPoolExecutor

# 初始化LMDB环境，设置足够大的映射空间（比如10GB）
env = lmdb.open("matrix_lmdb", map_size=1024*1024*1024*10)

# 封装写入逻辑：自动压缩
def write_matrix(key, matrix):
    compressed = blosc.compress(matrix.tobytes(), cname='lz4', clevel=5)
    with env.begin(write=True) as txn:
        txn.put(key.encode(), compressed)

# 封装读取逻辑：自动解压
def read_matrix(key):
    with env.begin() as txn:
        compressed = txn.get(key.encode())
        if compressed:
            data = blosc.decompress(compressed)
            return np.frombuffer(data, dtype=np.float32).reshape(224,224,1)

# 多线程写入：每个线程负责一批ID范围
def write_batch(start_id, end_id):
    for idx in range(start_id, end_id):
        mat = np.random.rand(224,224,1).astype(np.float32)
        write_matrix(f"matrix_{idx}", mat)

with ThreadPoolExecutor(max_workers=4) as executor:
    executor.submit(write_batch, 0, 250000)
    executor.submit(write_batch, 250000, 500000)
    executor.submit(write_batch, 500000, 750000)
    executor.submit(write_batch, 750000, 1000000)

4. 改进版HDF5（兼容现有习惯） #

如果你不想放弃HDF5，可以通过以下方式解决多线程写入问题：

• 分数据集写入：创建多个独立的数据集（比如每个线程写入一个专属数据集），或者按批次划分数据集，自行控制线程的写入范围，避免多个线程操作同一个数据集。
• 配合文件锁：使用portalocker等工具实现文件级锁，确保同一时间只有一个线程写入某个数据集，自己处理冲突逻辑。
• 透明压缩：HDF5本身支持gzip、lzf等压缩算法，创建数据集时配置即可。

简单使用示例：

import h5py
import numpy as np
from concurrent.futures import ThreadPoolExecutor
import portalocker

def write_dataset(ds_name, start_idx, end_idx):
    with h5py.File("matrix_hdf5.h5", "a") as f:
        # 获取文件锁，避免并发冲突
        portalocker.lock(f.id, portalocker.LOCK_EX)
        if ds_name not in f:
            f.create_dataset(
                ds_name,
                shape=(end_idx - start_idx, 224,224,1),
                dtype=np.float32,
                compression="gzip"
            )
        f[ds_name][:] = np.random.rand(end_idx - start_idx,224,224,1).astype(np.float32)
        portalocker.unlock(f.id)

with ThreadPoolExecutor(max_workers=4) as executor:
    executor.submit(write_dataset, "batch_0", 0, 250000)
    executor.submit(write_dataset, "batch_1", 250000, 500000)
    executor.submit(write_dataset, "batch_2", 500000, 750000)
    executor.submit(write_dataset, "batch_3", 750000, 1000000)

方案选择建议 #

• 如果以数组形式组织数据（按顺序或索引访问），优先选TileDB或Zarr+Dask，它们的API更贴合numpy数组习惯，并发支持更原生。
• 如果需要极致的读写性能，且习惯键值对存储，选LMDB，适合高频随机访问场景。
• 如果已有HDF5使用经验，不想切换技术栈，选改进版HDF5方案。

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