些細なことでも multiprocessing 使うと便利だよ、という小ネタ。
やりたいこと
from glob import glob
files = glob('data/*.csv')
len(files) # 10000
この 1万件の CSV ファイルを Pandas DataFrame として読み込みたい。
ちなみに検証用のデータは以下のようにして生成した。
(3列 x 10,000行 の CSV ファイル 10,000 個)
import numpy as np
import pandas as pd
row_n = 10000
col_n = 3
columns = [f'col{i}' for i in range(col_n)]
for i in range(10000):
df = pd.DataFrame(np.random.randn(row_n, col_n), columns=columns)
df.to_csv(f'data/{i:04}.csv', index=False)
並列化の手法いくつか
今回は
- 並列化なし (シングルスレッド)
- マルチスレッド
- マルチプロセス
で比較した。
asyncio (イベントループ) については、asyncio がファイルの非同期読み込みをサポートしていなかったので試さなかった。 aiofiles などの外部モジュールを使えばできると思うので、誰か試してみてほしい。 [追記] 試してみた → コメントを参照。
並列化なし (シングルスレッド)
並列化を考えずに全てのファイルを読み込もうとするとこうなる。
import pandas as pd
arr = [pd.read_csv(f) for f in files]
マルチプロセス
multiprocessing.Pool を使う。
import pandas as pd
from multiprocessing import Pool
def read_csv(f):
return pd.read_csv(f)
with Pool() as p:
arr = p.map(read_csv, files)
p.map には pd.read_csv を直接渡したり lambda を渡したりすることはできないので、 read_csv という関数を定義している。
Pool() で並列数を指定しなかった場合は CPU のコア数と同じ並列数で実行される。
マルチスレッド
multiprocessing.dummy.Pool は threading を使っているとのことなので、これを使ってみる。
import pandas as pd
from multiprocessing.dummy import Pool
with Pool() as p:
arr = p.map(pd.read_csv, files)
マルチスレッドの場合は pd.read_csv をそのまま渡せる。
結果
4コア CPU のマシンで計測した結果が以下。
| 読み込み速度 | |
|---|---|
| 並列化なし (シングルスレッド) | 82 sec |
| マルチスレッド | 28 sec |
| マルチプロセス | 22 sec |
マルチプロセスが一番速かった。
4並列で読み込んで4倍速くなるとまではいかなかったが、並列化の恩恵を十分に感じられる程度には速くなった。
マルチスレッドの注意点
今回はマルチスレッドもマルチプロセスとほぼ同じくらいには速くなったので、手軽さを重視してマルチスレッドの方を採用してもいいケースもあるかもしれない。
ただし Python (CPython) のマルチスレッドは GIL (Global Interpreter Lock) の制約があり、複数のスレッドが同時に Python バイトコードを実行することができない。
そのため、今回のファイル読み込みのような I/O バウンドな処理は高速化できるかもしれないが、例えば「読み込んだ各データに対してなんらかの重い処理を行う」などの CPU バウンドな処理を併せて実行した場合はおそらく今回ほど速くならない。そのようなケースではマルチプロセスで並列化した方が高速化できると思われる。
違う気がしてきた・・・。
詳しくはコメントを参照。
arr = [open(f).read() for f in files]
Pandas のパース処理を入れずに単純にファイル内容を読み込むだけにすると 3.2 秒で終わるので、今回の処理は I/O バウンドではなく CPU バウンドだった・・・。
[2023.02 追記] AWS Lambda で Pool が使えない件
AWS Lambda の Python ランタイムで同じことをやろうとしたら multiprocessing.Pool が使えなくて困ったので、それについて別記事に書いた。