我们在处理磁盘上的大型稀疏数据集的单细胞基因组数据领域遇到了类似的问题。我将向您展示一个简单的小例子，说明我将如何处理这个问题。我的假设是您的内存非常有限，并且可能无法一次将稀疏矩阵的多个副本放入内存中。即使您放不下一份完整的副本，这也可以使用。
我会逐列构造一个磁盘上的稀疏 CSC 矩阵。稀疏 csc 矩阵使用 3 个底层数组：
-
data：存储在矩阵中的值
-
indices: 矩阵中每个值的行索引
-
indptr：一个长度为n_cols + 1的数组，将indices和data除以它们所属的列。
作为一个解释性示例，i 列的值存储在data 的indptr[i]:indptr[i+1] 范围内。同样，这些值的行索引可以通过indices[indptr[i]:indptr[i+1]] 找到。
为了模拟您的数据生成过程（我假设是解析文档），我将定义一个函数 process_document，它返回相关文档的 indices 和 data 的值。
import numpy as np
import h5py
from scipy import sparse
from tqdm import tqdm # For monitoring the writing process
from typing import Tuple, Union # Just for argument annotation
def process_document():
"""
Simulate processing a document. Results in sparse vector represenation.
"""
n_items = np.random.negative_binomial(2, .0001)
indices = np.random.choice(2_000_000, n_items, replace=False)
indices.sort()
data = np.random.random(n_items).astype(np.float32)
return indices, data
def data_generator(n):
"""Iterator which yields simulated data."""
for i in range(n):
yield process_document()
现在我将在 hdf5 文件中创建一个组，该文件将存储稀疏矩阵的组成数组。
def make_sparse_csc_group(f: Union[h5py.File, h5py.Group], groupname: str, shape: Tuple[int, int]):
"""
Create a group in an hdf5 file that can store a CSC sparse matrix.
"""
g = f.create_group(groupname)
g.attrs["shape"] = shape
g.create_dataset("indices", shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,))
g["indptr"] = np.zeros(shape[1] + 1, dtype=int) # We want this to have a zero for the first value
g.create_dataset("data", shape=(1,), dtype=np.float32, chunks=True, maxshape=(None,))
return g
最后是一个将这个组读取为稀疏矩阵的函数（这个非常简单）。
def read_sparse_csc_group(g: Union[h5py.File, h5py.Group]):
return sparse.csc_matrix((g["data"], g["indices"], g["indptr"]), shape=g.attrs["shape"])
现在我们将创建磁盘上的稀疏矩阵并一次写入一列（我使用的列较少，因为这可能有点慢）。
N_COLS = 10
def make_disk_matrix(f, groupname, data_iter, shape):
group = make_sparse_csc_group(f, "mtx", shape)
indptr = group["indptr"]
data = group["data"]
indices = group["indices"]
n_total = 0
for doc_num, (cur_indices, cur_data) in enumerate(tqdm(data_iter)):
n_cur = len(cur_indices)
n_prev = n_total
n_total += n_cur
indices.resize((n_total,))
data.resize((n_total,))
indices[n_prev:] = cur_indices
data[n_prev:] = cur_data
indptr[doc_num+1] = n_total
# Writing
with h5py.File("data.h5", "w") as f:
make_disk_matrix(f, "mtx", data_generator(10), (2_000_000, 10))
# Reading
with h5py.File("data.h5", "r") as f:
mtx = read_sparse_csc_group(f["mtx"])
再次考虑到内存非常受限的情况，在这种情况下，您可能无法在创建时将整个稀疏矩阵放入内存中。如果您可以处理整个稀疏矩阵加上至少一个副本，那么执行此操作的一种更快的方法是不打扰磁盘存储（类似于其他建议）。但是，稍微修改一下这段代码应该会给你更好的性能：
def make_memory_mtx(data_iter, shape):
indices_list = []
data_list = []
indptr = np.zeros(shape[1]+1, dtype=int)
n_total = 0
for doc_num, (cur_indices, cur_data) in enumerate(data_iter):
n_cur = len(cur_indices)
n_prev = n_total
n_total += n_cur
indices_list.append(cur_indices)
data_list.append(cur_data)
indptr[doc_num+1] = n_total
indices = np.concatenate(indices_list)
data = np.concatenate(data_list)
return sparse.csc_matrix((data, indices, indptr), shape=shape)
mtx = make_memory_mtx(data_generator(10), shape=(2_000_000, 10))
这应该相当快，因为它只会在您连接数组后复制数据。当前发布的其他解决方案在您处理时重新分配了数组，从而制作了许多大型数组的副本。