temple

機械学習と周辺知識を整理する

Pandas でテーブルデータを読み書きする際のフォーマットごとの性能比較

動機

pandas を使うようになってからというもの、デファクトスタンダードである(と思っていた)csv形式でテーブルデータの読み書きをしていた。 最近は数GB程度のテーブルデータを扱う機会が増えてきており、データの読み書きをするだけでも相当な時間を要していることに気づいた。 「そもそもそんな規模感のデータを素の pandas で扱うな、daskを使え」などという批判も真摯に受け流しつつ、今一度 csv の代替となるフォーマットを探すべく我々はアマゾンの奥地へ向かった。

まずは結論

こまけえこたあいいんだよ

という人向けに。 結論、読み書き速度とデータサイズの観点だけで言えば csv からの乗り換え先は parquet 形式ではないでしょうか。 csv 形式と比較した結果は以下。

フォーマット データサイズ read 時間 write 時間
csv 3.5 GB 1 min 145 sec
parquet 1.1 GB 1.08 sec 19.3 sec

当然データを扱う状況によって最適なフォーマットは変わり得るので盲信はよくないが、「あんまりめんどくさいことはせずにとりあえず csv から乗り換えたい!」というユースケースには適合すると思う。

比較対象フォーマット

今回 csv 形式との比較対象としたものを以下表にまとめる。

形式 選定理由
jsonl 特に深い理由はない。強いて言うなら BigQuery のデータをエクスポートする際の選択肢として存在しているから…
parquet カラムナフォーマットの代表格。GCP の各種プロダクトでサポートしていることが多い。
pickle ML モデルを保存する際のデファクトスタンダード
joblib かつて個人的に pickle からの乗り換え先として愛用していたので。もはや pandas 関係ない。

上記のフォーマット以外にも様々な保存形式は存在するが、「極力 csv 形式からの引越しコストを抑えたい」「あまり馴染みのないフォーマットを扱いたくない」という個人的な怠慢から他のフォーマットは検討していない。 例えば TFRecord などは依存関係として tensorflow を持たせなきゃいけなかったり pandas とは異なるインターフェースでデータを扱う必要があり、データフォーマットを変えるだけでその副作用が出過ぎてしまうので候補に入れていない。

www.tensorflow.org

ちなみにフォーマットという観点以外にも cudf や dask を使うという選択肢もあるが、 cudf はまず GPU を用意するところから始まるので今までなんとな〜く csv を使っていた人がとりあえず使ってみるツールにしてはちょっとハードルが高い気がする。dask は pandas ライクに扱えるがインターフェースが違ったりハマりどころが色々あるので、しっかり腰を据えて高速化を検討するときの選択肢という印象…。

github.com

docs.dask.org

使用するデータ

以下のようにダミーデータを生成した。初めは調子に乗って 5,000 万行を float で作ろうとして普通に OOM で逝ったので控えめの integer 500万行にした。

import pandas as pd
import numpy as np

n_rows = 5_000_000
n_cols = 256
rng = np.random.default_rng()
table_df = pd.DataFrame(rng.integers(low=0, high=100, size=(n_rows, n_cols)))

path = "./table.{}"

write/read に使用するコード

保存フォーマットごとに使用したコードを記載する。

csv

# write
table_df.to_csv(path.format("csv"))
# read
table_df = pd.read_csv(path.format("csv"))

jsonl

# write
table_df.to_json(path.format("jsonl"), orient='records', lines=True)
# read
table_df = pd.read_json(path.format("jsonl"), orient='records', lines=True)

parquet

# write
table_df.to_parquet(path.format("parquet"))
# read
table_df = pd.read_parquet(path.format("parquet"))

pickle

# write
table_df.to_pickle(path.format("pkl"))
# read
table_df = pd.read_pickle(path.format("pkl"))

joblib

import joblib

# write
with open(path.format("joblib"), "wb") as f:
    joblib.dump(table_df, f, compress=3)
# read
with open(path.format("joblib"), "rb") as f:
    table_df = joblib.load(f)

(joblib だけ圧縮形式にしててズルいけど気にしない)

実行結果

jupyter notebook で%%timeマジックコマンドを使って実行した結果を表にまとめた。

フォーマット データサイズ read 時間 write 時間
csv 3.5 GB 1 min 145 sec
jsonl 11 GB N/A 177 sec
parquet 1.1 GB 1.08 sec 19.3 sec
pickle 9.6 GB 3.34 sec 42.3 sec
joblib 1.8 GB 25 sec 163 sec

jsonl は read しようとすると OOM で逝ってしまったので N/A とした。 表の通り、データサイズ・read 時間・write 時間の全項目で parquet 形式が圧勝だった。csv と比較して read は55倍、 write は7.5倍程度早くなっている。 ちなみに読み書きの速度に関しては pickle もいい線をいっているがデータサイズが csv の3倍弱になっているので、高速化の代償としては大きい気がする…。

結論

csv からなんか別なフォーマットにしたいな〜という人には parquet 形式をオススメします。 なお parquet はバイナリフォーマットなので、 csv のように vim とか head でファイルの中身をチラ見できないのがデメリットでしょうか。