Python awkward: parquet形式で1億イベントのtreeのIO(入出力)速度を圧縮形式で比較する

この記事の Awkward のバージョンは1.X.Xです

Python の awkward パッケージに実装されている parquet ファイル形式を入出力するときの処理時間を、圧縮形式ごとに比較してみる。配列の型は awkward.highlevel.Array と呼ばれる形式である。

8バイト(int64_t)✕8メンバー✕1億エントリーなので、圧縮なしの場合は6.4GBのデータサイズとなる。

圧縮形式 Compression algorithm	ファイルサイズ File size	出力時間 Write time	読み込み時間 Read time	理想IO時間 Ideal time
NONE(無圧縮)	6.40 GB	13.8 sec	16.5 sec	12.8 sec
SNAPPY(デフォルト)	3.21 GB	15.5 sec	18.4 sec	6.42 sec
LZ4	3.21 GB	11.0 sec	13.1 sec	6.42 sec
ZSTD	0.82 GB	19.0 sec	13.7 sec	1.64 sec
参考: ROOT LZ4	3.22 GB	37 sec	16 sec	6.44 sec
参考: ROOT ZSTD	0.56 GB	39 sec	25 sec	1.12 sec

理想IO時間は、SSDの読み書き速度を500MB/sとしたときのIO時間である。理想IO時間に、圧縮展開の時間、他の処理の時間を加えると、出力時間、読み込み時間になる。

NONEは、無圧縮なので6.4GBもの容量を使ってしまうものの、理想IO時間との差は当然ながら少なかった。デフォルト SNAPPY と LZ4 は 圧縮率は同じだが、IO速度はLZ4の圧勝だった。ZSTDは圧縮率は高いもののLZ4より遅い。LZ4かZSTDは、圧縮率を取るか、IO速度を取るかで選べばよさそうだ。

なお、C++でコンパイルしたROOTのTFileのIO速度と同じアルゴリズムで比べると(圧縮率の差はあるが)Pythonの方が2倍早い。Treeをシンプルな処理で使う分にはPythonで十分であることが分かるだろう。

Python awkward の parquet ファイルのIO速度を検証したコード。ディスクキャッシュを防ぐため、数20GBのディスク容量を使う。

import awkward as ak
assert(ak.__version__.split(".")[0]=="1")
import numpy as np
import time
import os

comps = ["NONE", "SNAPPY", "LZ4", "ZSTD"] #"GZIP", "BROTLI"は遅すぎるので除外
for comp in comps:
    for i in range(3):
        start = time.time()

        N = 1 * 10000 * 10000
        tree = ak.Array({})
        tree["x"] = np.arange(N, dtype=np.int64)
        tree["y"] = np.arange(N, dtype=np.int64)
        tree["ax"] = np.arange(N, dtype=np.int64)
        tree["ay"] = np.arange(N, dtype=np.int64)
        tree["ax0"] = np.arange(N, dtype=np.int64)
        tree["ay0"] = np.arange(N, dtype=np.int64)
        tree["ax1"] = np.arange(N, dtype=np.int64)
        tree["ay1"] = np.arange(N, dtype=np.int64)
        
        filename = f"tree{i}_{comp}.parquet"
        ak.to_parquet(tree, filename, compression=comp)

        end = time.time()
        print(f"WRITE {i} {comp} {os.path.getsize(filename)/1e9:.2f} GB {(end-start):.1f} sec")

        del tree

import awkward as ak
import numpy as np
import time

comps = ["NONE", "SNAPPY", "LZ4", "ZSTD"] #"GZIP", "BROTLI"は遅すぎるので除外

for comp in comps:
    for i in range(3):
        start = time.time()
        filename = f"tree{i}_{comp}.parquet"
        tree = ak.from_parquet(filename, lazy=True)
        for key in ["x","y","ax","ay","ax0","ay0","ax1","ay1"]:
            print(np.max(tree[key]),end=" ")
        end = time.time()
        print(f"READ {i} {comp} {(end-start):.1f} sec")
        
        del tree

Awkward 2.X.Xを使う場合は、おそらく以下のように書き換えるとよい

import awkward as ak
assert(ak.__version__.split(".")[0]=="2")
import numpy as np
import time
import os

comps = ["NONE", "SNAPPY", "LZ4", "ZSTD"] #"GZIP", "BROTLI"は遅すぎるので除外
for comp in comps:
    for i in range(3):
        start = time.time()

        N = 1 * 10000 * 10000
        tree = ak.Array({"x":np.arange(N, dtype=np.int64)})
        for key in ["y","ax","ay","ax0","ay0","ax1","ay1"]:
            tree[key] = np.arange(N, dtype=np.int64)
        
        filename = f"tree{i}_{comp}.parquet"
        ak.to_parquet(tree, filename, compression=comp)

        end = time.time()
        print(f"WRITE {i} {comp} {os.path.getsize(filename)/1e9:.2f} GB {(end-start):.1f} sec")

        del tree

import dask_awkward as dak
import awkward as ak
import time

comps = ["NONE", "SNAPPY", "LZ4", "ZSTD"] #"GZIP", "BROTLI"は遅すぎるので除外

for comp in comps:
    for i in range(3):
        start = time.time()
        filename = f"tree{i}_{comp}.parquet"
        tree = dak.from_parquet(filename)
        for key in ["x","y","ax","ay","ax0","ay0","ax1","ay1"]:
            print(ak.max(tree[key]).compute(),end=" ")
        end = time.time()
        print(f"READ {i} {comp} {(end-start):.1f} sec")
        
        del tree