如何将大型稀疏Dask数组转换为单个SciPy CSR矩阵？

我有一个20万行、15万列的稀疏Dask数组，需要转换成单个SciPy稀疏CSR矩阵，用于implicit库的协同过滤模型。

当前尝试的思路 #

思路1：直接使用compute方法 #

尝试直接计算并转换，但因内存不足崩溃：

import dask.array as da
import sparse
nusers = 200_000
nitems = 150_000
rng = da.random.default_rng()
x = rng.random((nusers, nitems))
x[x < 0.95] = 0
s = x.map_blocks(sparse.COO)
# 此操作因内存不足导致程序崩溃
s_csr = s.compute().tocsr()

思路2：增量创建稀疏数组并写入磁盘 #

尝试分块拼接后转CSR写入磁盘，但速度极慢，还需手动GC：
注意：sparse.concatenate()得到的矩阵占用内存会持续增长，但转成SciPy CSR后仅占48字节（测试得出）。

import sys
import gc
from scipy.sparse import save_npz, csr_matrix
import dask.array as da
import sparse

def get_next_chunk(s: da.Array):
    """获取DaskArray的下一个块""" 
    raise NotImplementedError

# 重新分块，沿第一轴增量拼接
s = s.rechunk((some_value, nitems))

# 获取第一个块
sp_matrix = get_next_chunk(s)

# 遍历剩余块拼接
while (chunk := get_next_chunk(s)):
    sp_matrix = sparse.concatenate((sp_matrix, chunk))
    # 内存阈值触发写入
    if (sys.getsizeof(sp_matrix)/2**30) > threshold_GB:
        sp_matrix = sp_matrix.tocsr() 
        save_npz(fpath, sp_matrix)
        sp_matrix = sparse.COO.from_scipy_sparse(sp_matrix)
        gc.collect()

sp_matrix = sp_matrix.tocsr()
save_npz(final, sp_matrix)

硬件环境 #

搭载8GB内存和8GB交换空间的Lubuntu 22.04 LTS笔记本电脑

针对内存限制和需求，推荐直接操作CSR矩阵的底层数组进行增量构建，避免中间大对象占用内存，同时提升处理速度：

核心思路 #

SciPy CSR矩阵由indptr（行指针）、indices（列索引）、data（非零值）三个核心数组构成。我们可以逐块处理Dask数组的行块，直接累加这三个数组，最后拼接成完整的CSR矩阵，全程跳过高开销的稀疏对象拼接操作。

具体实现代码 #

import dask.array as da
import sparse
from scipy.sparse import csr_matrix
import numpy as np
import gc

# 假设你的Dask稀疏数组已经定义为s
nusers = 200_000
nitems = 150_000

# 调整分块大小，根据内存情况选择（比如1000行/块，内存紧张可减小到500）
chunk_size = 1000
s = s.rechunk((chunk_size, nitems))

# 初始化CSR核心数组
indptr = [0]
indices = []
data = []

# 逐块处理
for block in s.blocks:
    # 计算当前块的COO矩阵并转CSR
    coo_block = block.compute()
    csr_block = coo_block.tocsr()
    
    # 累加行指针：当前块的行指针偏移之前的总非零数
    indptr.extend(csr_block.indptr[1:] + indptr[-1])
    # 累加列索引和非零值
    indices.extend(csr_block.indices)
    data.extend(csr_block.data)
    
    # 手动释放临时对象，避免内存堆积
    del coo_block, csr_block
    gc.collect()

# 转换为numpy数组并构建完整CSR矩阵
indptr = np.array(indptr, dtype=np.int64)
indices = np.array(indices, dtype=np.int32)
data = np.array(data, dtype=np.float64)

full_csr = csr_matrix((data, indices, indptr), shape=(nusers, nitems))

优化说明 #

1. 内存控制：每次仅处理一个小块，仅保留三个核心数组，内存占用仅为所有非零元素的存储大小，完全适配你的硬件。
2. 效率提升：跳过sparse.concatenate的高开销操作，直接操作CSR底层数组，处理速度比增量拼接快数倍。
3. 内存泄漏预防：每处理完一块就手动释放临时对象并触发GC，避免内存溢出。

额外建议 #

• 调整chunk_size：如果内存仍紧张，可减小分块大小；若有剩余内存，可增大分块以提升处理效率。
• 数据类型优化：如果业务允许，将data改为float32、indices改为int32，进一步降低内存占用。
• 监控交换空间：若处理中频繁触发交换，优先减小分块大小，避免磁盘IO拖慢速度。

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