Легковесный и сверхбыстрый I/O для массивов NumPy на LZ4

Всем привет! 👋

При обработке больших объемов числовых данных мы постоянно сталкиваемся с классическим компромиссом между скоростью дискового ввода-вывода и занимаемым местом. Стандартная экосистема NumPy вынуждает идти на неприятные уступки при простом кэшировании промежуточных массивов:

- np.save() работает молниеносно, но не дает никакого сжатия, впустую расходуя место на диске и пропускную способность сети.

- np.savez_compressed() (zlib) экономит место, но сильно нагружает процессор (CPU-bound). Запись всего лишь 100 МБ массива может занимать несколько секунд, что неприемлемо для циклов обучения моделей или систем реального времени.

- HDF5 / Zarr — невероятно мощные инструменты, но тащить тяжелые зависимости и настраивать чанкинг (chunking) — это избыточно, если нужно просто сохранить временное состояние.

Чтобы решить эту проблему, я реализовал легковесное решение, которое объединяет скорость прямого доступа к памяти и эффективность алгоритма LZ4 через C-расширение для Python.

Как это работает:

Критический путь сериализации вынесен в C, что позволяет обойти накладные расходы интерпретатора Python и глобальную блокировку (GIL). Используется Buffer Protocol (PyArray_DATA) для прямого доступа к памяти, что обеспечивает подход, близкий к zero-copy. Данные сериализуются со строгим кроссплатформенным бинарным заголовком размером 96 байт (включая магическое число для мгновенной валидации). Это гарантирует безопасную переносимость между архитектурами (например, x86_64 и ARM64) и предотвращает попытки распаковки поврежденных файлов.

Результаты:

В бенчмарках на массивах float32 объемом ~1 МБ подход на базе LZ4 показал задержку записи, сопоставимую с несжатым no.save (доли миллисекунды), при этом оказавшись примерно в 10 раз быстрее zlib.

Что, возможно, наиболее неожиданно: на числовых данных LZ4 в этих тестах стабильно показывал лучшую степень сжатия, чем zlib, уменьшая тестовый массив с 1 МБ до ~0.4 МБ (против ~0.5 МБ у zlib). Это достигается благодаря особенностям работы со словарем и отсутствию избыточных проверок, присущих алгоритму deflate, а также зависит от параметра acceleration, который указывает LZ4 с какой степенью сжимать данные.

API намеренно сделано минималистичным (всего две функции save и load, которые автоматически обрабатывают метаданные dtype и shape), что делает библиотеку бесшовной заменой (drop-in replacement) без лишней головной боли. Проект распространяется с открытым исходным кодом под лицензией Apache 2.0. 🔗 PyPI: https://pypi.org/project/numpy-cache/ 🔗 GitHub: https://github.com/macht1212/numpy-cache