So I don't know anything about joblib or dask, let alone your application specific data format. But it is actually possible to read sparse matrices from disk in chunks while retaining the sparse data structure.
While the Wikipedia article for the CSR format does a great job explaining how it works, I'll give a short recap:
Some sparse Matrix, e.g.:
1 0 2
0 0 3
4 5 6
is stored by remembering each nonzero-value and the column it resides in:
sparse.data = 1 2 3 4 5 6 # acutal value
sparse.indices = 0 2 2 0 1 2 # number of column (0-indexed)
Now we are still missing the rows. The compressed format just stores how many non-zero values there are in each row, instead of storing every single values row.
Note that the non-zero count is also accumulated, so the following array contains the number of non-zero values up until and including this row. To complicate things even further, the array always starts with a 0
and thus contains num_rows+1
entries:
sparse.indptr = 0 2 3 6
so up until and including the second row there are 3 nonzero values, namely 1
, 2
and 3
.
Since we got this sorted out, we can start 'slicing' the matrix. The goal is to construct the data
, indices
and indptr
arrays for some chunks. Assume the original huge matrix is stored in three binary files, which we will incrementally read. We use a generator to repeatedly yield
some chunk.
For this we need to know how many non-zero values are in each chunk, and read the according amount of values and column-indices. The non-zero count can be conveniently read from the indptr array. This is achieved by reading some amount of entries from the huge indptr
file that corresponds to the desired chunk size. The last entry of that portion of the indptr
file minus the number of non-zero values before gives the number of non-zeros in that chunk. So the chunks data
and indices
arrays are just sliced from the big data
and indices
files. The indptr
array needs to be prepended artificially with a zero (that's what the format wants, don't ask me :D).
Then we can just construct a sparse matrix with the chunk data
, indices
and indptr
to get a new sparse matrix.
It has to be noted that the actual matrix size cannot be directly reconstructed from the three arrays alone. It is either the maximum column index of the matrix, or if you are unlucky and there is no data in the chunk undetermined. So we also need to pass the column count.
I probably explained things in a rather complicated way, so just read this just as opaque piece of code, that implements such a generator:
import numpy as np
import scipy.sparse
def gen_batches(batch_size, sparse_data_path, sparse_indices_path,
sparse_indptr_path, dtype=np.float32, column_size=None):
data_item_size = dtype().itemsize
with open(sparse_data_path, 'rb') as data_file,
open(sparse_indices_path, 'rb') as indices_file,
open(sparse_indptr_path, 'rb') as indptr_file:
nnz_before = np.fromstring(indptr_file.read(4), dtype=np.int32)
while True:
indptr_batch = np.frombuffer(nnz_before.tobytes() +
indptr_file.read(4*batch_size), dtype=np.int32)
if len(indptr_batch) == 1:
break
batch_indptr = indptr_batch - nnz_before
nnz_before = indptr_batch[-1]
batch_nnz = np.asscalar(batch_indptr[-1])
batch_data = np.frombuffer(data_file.read(
data_item_size * batch_nnz), dtype=dtype)
batch_indices = np.frombuffer(indices_file.read(
4 * batch_nnz), dtype=np.int32)
dimensions = (len(indptr_batch)-1, column_size)
matrix = scipy.sparse.csr_matrix((batch_data,
batch_indices, batch_indptr), shape=dimensions)
yield matrix
if __name__ == '__main__':
sparse = scipy.sparse.random(5, 4, density=0.1, format='csr', dtype=np.float32)
sparse.data.tofile('sparse.data') # dtype as specified above ^^^^^^^^^^
sparse.indices.tofile('sparse.indices') # dtype=int32
sparse.indptr.tofile('sparse.indptr') # dtype=int32
print(sparse.toarray())
print('========')
for batch in gen_batches(2, 'sparse.data', 'sparse.indices',
'sparse.indptr', column_size=4):
print(batch.toarray())
the numpy.ndarray.tofile()
just stores binary arrays, so you need to remember the data format. scipy.sparse
represents the indices
and indptr
as int32
, so that's a limitation for the total matrix size.
Also I benchmarked the code and found that the scipy csr matrix constructor is the bottleneck for small matrices. Your mileage might vary tho, this is just a 'proof of principle'.
If there is need for a more sophisticated implementation, or something is too obstruse, just hit me up :)