CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: pickle を使って pandas の CSV 読み込み時間を削減する

機械学習やデータ分析に使うデータセットは CSV などの形式で提供される場合がある。 それを Python で処理するときは pandas の DataFrame 形式に変換することが多い。 このとき CSV から DataFrame に変換する処理は意外と時間がかかる。 特に大きなデータセットでこれを毎回やっていると効率が悪い。 今回は、一旦読み込んだ DataFrame を pickle を使って直接ファイルに保存することで時間を節約できるという話について。

使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.13.4
BuildVersion:   17E202
$ python -V                                
Python 3.6.5

下準備

まずは pandas と iptyhon をインストールしておく。 ipython を入れているのは、実行にかかる時間を測るのに便利なため。

$ pip install pandas ipython

サンプルとして大きな CSV のデータセットが欲しかったので適当に Kaggle から見繕ってきた。 これは大気の汚染状況に関するデータセット。

$ kaggle datasets download -d epa/hazardous-air-pollutants

kaggle コマンドについては以下の記事を参照のこと。

blog.amedama.jp

もちろん Kaggle の Web サイトからダウンロードしてきても問題はない。

Hazardous Air Pollutants | Kaggle

ダウンロードしてきたデータセットを解凍する。

$ mv ~/.kaggle/datasets/epa/hazardous-air-pollutants/hazardous-air-pollutants.zip .
$ unzip hazardous-air-pollutants.zip 
Archive:  hazardous-air-pollutants.zip
  inflating: epa_hap_daily_summary.csv

これで 2GB 強の CSV ファイルが表れる。

$ du -m epa_hap_daily_summary.csv 
2350   epa_hap_daily_summary.csv

CSV を DataFrame に変換する

それでは上記の CSV を DataFrame に変換してみよう。

まずは IPython を起動する。

$ ipython

起動したら pandas をインポートしよう。

>>> import pandas as pd

あとは pandas.read_csv() 関数で CSV を DataFrame に変換する。 このとき IPython の %time マジックコマンドで処理にかかる時間を測っておこう。

>>> %time df = pd.read_csv('epa_hap_daily_summary.csv')
CPU times: user 55 s, sys: 9.76 s, total: 1min 4s
Wall time: 1min 7s

読み込みには 1 分強かかった。

読み込まれた DataFrame は次の通り。

>>> df.head()
state_code  county_code  site_num  parameter_code  poc   latitude  \
0          42          125      5001           88103    5  40.445278   
1          48          439      1002           43843   15  32.805818   
2          22          127      9000           88128    1  32.057581   
3          18           89        22           45201    1  41.606680   
4           6           89      3003           88136    1  40.539990   

 longitude  datum      parameter_name sample_duration         ...           \
0  -80.420833  WGS84    Arsenic PM2.5 LC         24 HOUR         ...            
1  -97.356568  WGS84  Ethylene dibromide         24 HOUR         ...            
2  -92.435157  WGS84       Lead PM2.5 LC         24 HOUR         ...            
3  -87.304729  WGS84             Benzene          1 HOUR         ...            
4 -121.576460  NAD83     Nickel PM2.5 LC         24 HOUR         ...            
...(省略)

さて、今回かかった時間は 1 分とはいえちりも積もればという話もある。 また、より複雑なパースを指定すれば必要な時間はどんどん増えていく。

DataFrame を pickle で保存・復元する

そこで、次は DataFrame を pickle で保存・復元してみる。 CSV から逐一変換するのに比べて、どれくらい速くなるだろうか。

pandas の DataFrame には to_pickle() というメソッドがあるので、それを使えば DataFrame をファイルに保存できる。

>>> df.to_pickle('hazardous-air-pollutants.pickle')

約 1.5GB の pickle ファイルができた。

$ du -m hazardous-air-pollutants.pickle 
1501   hazardous-air-pollutants.pickle

それでは、上記のファイルを改めて読み込んでみる。 そのために、ここで一旦インタプリタを終了しておこう。

>>> exit()

もう一度 IPython を起動し直す。

$ ipython

まずは pickle モジュールをインポートする。

>>> import pickle

そして、次のようにして保存したファイルから DataFrame を復元する。 こちらも、処理に要する時間を %%time マジックコマンドで計測しておこう。

>>> %%time
... with open('hazardous-air-pollutants.pickle', mode='rb') as fp:
...     df = pickle.load(fp)
...
CPU times: user 4.5 s, sys: 2.24 s, total: 6.74 s
Wall time: 7.21 s

今度は約 7 秒で処理が終わった。 ここまで速くなる理由は CSV の各データを識別・パースする処理が不要になるため。

あるいは pandas の API を使って次のように一行で読み込むこともできる。

>>> %time df = pd.read_pickle('hazardous-air-pollutants.pickle')
CPU times: user 4.51 s, sys: 4.07 s, total: 8.58 s
Wall time: 9.03 s

ちゃんと DataFrame が復元されている。

>>> df.head()
state_code  county_code  site_num  parameter_code  poc   latitude  \
0          42          125      5001           88103    5  40.445278   
1          48          439      1002           43843   15  32.805818   
2          22          127      9000           88128    1  32.057581   
3          18           89        22           45201    1  41.606680   
4           6           89      3003           88136    1  40.539990   

 longitude  datum      parameter_name sample_duration         ...           \
0  -80.420833  WGS84    Arsenic PM2.5 LC         24 HOUR         ...            
1  -97.356568  WGS84  Ethylene dibromide         24 HOUR         ...            
2  -92.435157  WGS84       Lead PM2.5 LC         24 HOUR         ...            
3  -87.304729  WGS84             Benzene          1 HOUR         ...            
4 -121.576460  NAD83     Nickel PM2.5 LC         24 HOUR         ...            

めでたしめでたし。