Внешняя обработка разреженных массивов CSR

Так что я ничего не знаю о joblib или dask, не говоря уже о формате данных вашего приложения. Но на самом деле можно читать разреженные матрицы с диска порциями, сохраняя при этом разреженную структуру данных.

В то время как статья Википедии о формате CSR отлично объясняет как это работает, я дам краткий обзор:

Некоторая разреженная матрица, например:

1 0 2
0 0 3
4 5 6

сохраняется путем запоминания каждого ненулевого значения и столбца, в котором оно находится:

sparse.data    = 1 2 3 4 5 6  # acutal value
sparse.indices = 0 2 2 0 1 2  # number of column (0-indexed)

Теперь нам все еще не хватает строк. В сжатом формате просто сохраняется количество ненулевых значений в каждой строке, а не каждая отдельная строка значений.

Обратите внимание, что ненулевое количество также накапливается, поэтому следующий массив содержит количество ненулевых значений до этой строки включительно. Чтобы еще больше усложнить ситуацию, массив всегда начинается с 0 и, таким образом, содержит num_rows+1 записей:

sparse.indptr = 0 2 3 6

поэтому до второй строки включительно есть 3 ненулевых значения, а именно 1, 2 и 3.

Поскольку мы с этим разобрались, мы можем начать «нарезать» матрицу. Цель состоит в том, чтобы построить массивы data, indices и indptr для некоторых фрагментов. Предположим, исходная огромная матрица хранится в трех бинарных файлах, которые мы будем постепенно читать. Мы используем генератор для повторного yield некоторого фрагмента.

Для этого нам нужно знать, сколько ненулевых значений в каждом фрагменте, и прочитать соответствующее количество значений и индексов столбцов. Ненулевой счетчик удобно считывать из массива indptr. Это достигается чтением некоторого количества записей из огромного файла indptr, соответствующего желаемому размеру фрагмента. Последняя запись этой части файла indptr за вычетом количества ненулевых значений перед дает количество ненулевых значений в этом фрагменте. Таким образом, массивы фрагментов data и indices просто вырезаны из больших файлов data и indices. Массив indptr должен быть искусственно добавлен нулем (это то, чего хочет формат, не спрашивайте меня: D).

Затем мы можем просто построить разреженную матрицу с фрагментами data, indices и indptr, чтобы получить новую разреженную матрицу.

Следует отметить, что фактический размер матрицы не может быть напрямую восстановлен только по трем массивам. Это либо максимальный индекс столбца матрицы, либо если вам не повезло и в чанке нет неопределенных данных. Поэтому нам также нужно передать количество столбцов.

Я, наверное, довольно сложно объяснил, так что просто прочитайте это просто как непрозрачный кусок кода, который реализует такой генератор:

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())

numpy.ndarray.tofile() просто хранит двоичные массивы, поэтому вам нужно запомнить формат данных. scipy.sparse представляет indices и indptr как int32, так что это ограничение общего размера матрицы.

Также я проверил код и обнаружил, что конструктор матрицы scipy csr является узким местом для небольших матриц. Ваш пробег может отличаться, это просто «доказательство принципа».

Если есть необходимость в более сложной реализации или что-то слишком запутанное, просто напишите мне :)