Что такое интеллектуальный способ полной загрузки сжатого массива с диска в память - также (по сути) сжатого?

Я экспериментирую с трехмерным zarr-массивом, хранящимся на диске:

Name: /data
Type: zarr.core.Array
Data type: int16
Shape: (102174, 1100, 900)
Chunk shape: (12, 220, 180)
Order: C
Read-only: True
Compressor: Blosc(cname='zstd', clevel=3, shuffle=BITSHUFFLE, blocksize=0)
Store type: zarr.storage.DirectoryStore
No. bytes: 202304520000 (188.4G)
No. bytes stored: 12224487305 (11.4G)
Storage ratio: 16.5
Chunks initialized: 212875/212875

Насколько я понимаю, zarr-массивы также могут находиться в памяти - сжаты, как если бы они были на диске. Поэтому я подумал, почему бы не попробовать загрузить все это в оперативную память на машине с 32 ГБ памяти. Сжатый набор данных потребует около 50% оперативной памяти. Несжатый, он потребует примерно в 6 раз больше оперативной памяти, чем доступно.

Приготовление:

import os
import zarr
from numcodecs import Blosc
import tqdm
zpath = '...' # path to zarr data folder

disk_array = zarr.open(zpath, mode = 'r')['data']

c = Blosc(cname = 'zstd', clevel=3, shuffle = Blosc.BITSHUFFLE)
memory_array = zarr.zeros(
    disk_array.shape, chunks = disk_array.chunks,
    dtype = disk_array.dtype, compressor = c
    )

Следующий эксперимент завершается почти сразу же с ошибкой нехватки памяти:

memory_array[:, :, :] = disk_array[:, :, :]

Как я понимаю, disk_array[:, :, :] попытается создать несжатый полноразмерный массив numpy, который, очевидно, потерпит неудачу.

Вторая попытка, которая работает, но мучительно медленно:

chunk_lines = disk_array.chunks[0]
chunk_number = disk_array.shape[0] // disk_array.chunks[0]
chunk_remain = disk_array.shape[0] % disk_array.chunks[0] # unhandled ...
for chunk in tqdm.trange(chunk_number):
    chunk_slice = slice(chunk * chunk_lines, (chunk + 1) * chunk_lines)
    memory_array[chunk_slice, :, :] = disk_array[chunk_slice, :, :]

Здесь я пытаюсь прочитать определенное количество фрагментов за раз и поместить их в мой массив в памяти. Это работает, но это примерно в 6-7 раз медленнее, чем то, что потребовалось для записи этой вещи на диск. РЕДАКТИРОВАТЬ: Да, это все еще медленно, но 6-7 раз произошло из-за проблемы с диском.

Какой разумный и быстрый способ достижения этого? Я предполагаю, что помимо неправильного подхода мои куски могут быть слишком маленькими, но я не уверен.

РЕДАКТИРОВАТЬ: форма, размер куска и сжатие должны быть идентичны для массива на диске и массива в памяти. Следовательно, должна быть возможность исключить процедуру распаковки-сжатия в моем примере выше.

я нашел zarr.convenience.copy но это помечено как experimental feature, подлежит дальнейшему изменению.


Связанная проблема на GitHub

2 ответа

Есть несколько способов решить эту проблему сегодня.

  1. использование LRUStoreCache кэшировать (некоторые) сжатые данные в памяти.
  2. Приведите ваш основной магазин в dict и использовать это как свой магазин.

Первый вариант может быть уместным, если вам нужны только некоторые часто используемые данные в памяти. Конечно, сколько вы загружаете в память, это то, что вы можете настроить. Так что это может быть весь массив. Это произойдет только с данными по запросу, которые могут быть полезны для вас.

Второй вариант просто создает новую копию массива в памяти путем извлечения всех сжатых данных с диска. Единственным недостатком является то, что если вы намереваетесь записать обратно на диск, это будет то, что вам нужно сделать вручную, но это не слишком сложно. update Метод очень удобен для облегчения копирования данных между различными хранилищами.

Вы могли бы попытаться с fsspec.implementations.memory.MemoryFileSystem, который имеет .make_mapper() метод, с помощью которого вы можете сделать вид объекта, ожидаемый Zarr.

Тем не менее, это действительно просто путь:io.BytesIO, который вы можете сделать сами, если хотите.

Другие вопросы по тегам