🤐

Blosc2 で Numpy Array を高速に圧縮・展開する

2024/09/15に公開

numpy

Numpy Array を保存するとき、numpy に用意されている np.savez_compressed を使うのが一般的だと思いますが、blosc2.pack_array2/blosc2.unpack_array2 が速度・圧縮率の両面で優秀なので、適当なデータで圧縮してみた実験結果を紹介したいと思います。

最も典型的な Dense なデータで比較した結果だけ先に紹介しておくと、np.savez_compressed に比べて、圧縮が38倍、展開が16倍高速で、ファイルサイズも 8% 程度縮みます。

インストール

普通に pip でインストールできます

pip install blosc2

使用方法

import blosc2
import numpy as np


def save_blosc2(path: str, x: np.ndarray) -> None:
    with open(path, "wb") as f:
        f.write(blosc2.pack_array2(x))


def load_blosc2(path: str) -> np.ndarray:
    with open(path, "rb") as f:
        return blosc2.unpack_array2(f.read())

blosc2.pack_array/unpack_array (2 がついていないもの) もありますが、これは blosc 無印時代への backward competibility のために残っている古い関数です。2 がついている方を使うように注意してください。

実験設定

比較対象1: `np.savez_compressed`

def save_npz(path: str, x: np.ndarray) -> None:
    np.savez_compressed(path, x)


def load_npz(path: str) -> np.ndarray:
    return np.load(path)["arr_0"]

比較対象2: `zstandard`

pip install zstandard

zstandard ライブラリには Numpy の Array を扱うような関数は無いため、np.load, np.save と組み合わせて利用します。

import zstandard


def save_zstd(path: str, x: np.ndarray) -> None:
    with zstandard.open(path, "wb") as f:
        np.save(f, x)


def load_zstd(path) -> np.ndarray:
    with zstandard.open(path, "rb") as f:
        return np.load(io.BytesIO(f.read()))

データ

以下の傾向の異なる 3 つのデータで実験しました。

Dense: ランダムな fp32
Sparse: 0.01% だけ 1 で他はすべて 0
Linear: np.arange

shape = (256, 3, 256, 256)  # 適当な shape

x_dense = np.random.RandomState(0).uniform(0, 1, size=shape).astype(np.float32)
print(f"{x_dense.nbytes / 1024 ** 2} MiB")
# 192.0 MiB

x_sparse = np.zeros(shape, dtype=np.float32)
x_sparse[np.unravel_index(np.random.RandomState(0).permutation(x_sparse.size)[:int(x_sparse.size * 0.01 / 100)], shape)] = 1
print(f"{x_sparse.mean() * 100:.2g}%, {x_sparse.nbytes / 1024 ** 2} MiB")
# 0.01%, 192.0 MiB

x_linear = np.arange(x_sparse.size).reshape(shape).astype(np.int32)
print(f"{x_linear.dtype}, {x_linear.nbytes / 1024 ** 2} MiB")
# int32, 192.0 MiB

その他

実行環境は MacBook Air の M2, 2022 で、Python 3.11.8 と以下のバージョンの package を使用しています。

blosc2==2.7.1
numpy==2.1.1
zstandard==0.23.0

実験コードは全体はこちら

結果

圧縮時間

	dense	sparse	linear
npz	6.05	0.753	14.3
zstandard	0.257	0.0321	0.281
blosc2	0.156	0.00823	0.00931

爆速！！！

展開時間

	dense	sparse	linear
npz	0.764	0.3	0.494
zstandard	0.252	0.0749	0.202
blosc2	0.047	0.0161	0.012

爆速!!!

ファイルサイズ

	dense	sparse	linear
npz	172	0.205	66.4
zstandard	172	0.0536	130
blosc2	157	0.126	0.823

(Sparseは zstandard に負けているけど) 小さい!!!

まとめ

blosc2.pack_array2/unpack_array2 が速度・圧縮率の両面で優秀ですね！

今回はデフォルトの設定でのみ比較を行いましたが、blosc2 には BLOSCLZ, LZ4, LZ4HC, ZLIB, ZSTD などの様々な Codec が実装されていて、pack_array2 の cparams に以下のように指定することで Codec を指定することができます。用途に合わせて Codec を指定したり圧縮の設定を変更することでさらなる高速化が望めるかもしれません。

blosc2.pack_array2(x, cparams={"codec": blosc2.Codec.BLOSCLZ})

zstandard との比較でも、Sparse なデータのファイルサイズを除けば blosc2 の方が優れた結果となっていましたが、実は blosc2 のデフォルトの圧縮の設定では　Codec としては BLOSCLZ ではなく ZSTD が使われているよう (該当部分のコード) なので、Numpy に特化した最適化が入っているかどうかが zstandard との差を生んでいるのだと考えています (詳しいところは確認していません！)。

今回は blosc2 の pack_array2/unpack_array2 のみ取り上げましたが、blosc2 のメインの機能は NDarray のような Numpy Array を圧縮したまま保持できるようなデータ構造で、Tutorialではそこら辺が紹介されているので、興味のある方は読んでみると面白いと思います。