【Mac OS】巨大なデータをpickleで読み書きしたときにはまったこと

あまりデータサイズが大きくない訓練データセットを読み込む時は普通にpandasのread_csvを使って

import pandas as pd
df = pd.read_csv('train.csv')

で大丈夫だけど、pandasのread_csvは遅いので、データサイズが大きい時は他の方法を使いたくなる。 (実行環境:MacBook Pro (Early 2015)、CPU 2.9 GHz Intel Core i5、メモリ 8 GB 1867 MHz DDR3、データCSV:7.0GB)

%time df = pd.read_csv('train.csv')
CPU times: user 2min 31s, sys: 58.8 s, total: 3min 29s
Wall time: 4min 2s

遅い。。より良い方法がないか、と調べているとpickleでシリアライズ読み書きすれば良いとわかるが、以下のコードだとうまくいかない。

import pickle
pickle.dump(df, open('train.csv.pkl', 'wb'), protocol=4)

OSError Traceback (most recent call last)
<ipython-input-10-aa4f708bf30d> in <module>()
 1 import pickle
----> 2 pickle.dump(df, open('train.csv.pkl', 'wb'), protocol=4)

OSError: [Errno 22] Invalid argument

このエラーをStackover flowで調べてみると長い間解決されていない、Mac OSでは2GBより大きいデータ書き込み時にエラーが発生する問題が原因だということが分かる。 手っ取り早い解決策として、以下のコードのように2の31乗バイトに分けてファイル出力するとエラーがでなくなるらしい。

import sys
import os

def save_as_pickled_object(obj, filepath):
    max_bytes = 2**31 - 1
    bytes_out = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
    n_bytes = sys.getsizeof(bytes_out)
    with open(filepath, 'wb') as f_out:
        for idx in range(0, n_bytes, max_bytes):
            f_out.write(bytes_out[idx:idx+max_bytes])


def try_to_load_as_pickled_object_or_None(filepath):
    max_bytes = 2**31 - 1
    try:
        input_size = os.path.getsize(filepath)
        bytes_in = bytearray(0)
        with open(filepath, 'rb') as f_in:
            for _ in range(0, input_size, max_bytes):
                bytes_in += f_in.read(max_bytes)
        obj = pickle.loads(bytes_in)
    except:
        return None
    return obj
# 書き込み
save_as_pickled_object(df, 'train.csv.pkl')

# 読み込み
%time df = try_to_load_as_pickled_object_or_None('train.csv.pkl')
CPU times: user 27.8 s, sys: 1min 26s, total: 1min 53s
Wall time: 3min 7s

pickleで読み書きできるようにはなったが、結局1分くらいしか早くなっていないので、pickleからcPickle(python 3.xでbundleされているpickleのC言語実装)にしたところ少し早くなった。

import _pickle as cPickle
import sys
import os

def try_to_load_as_pickled_object_or_None(filepath):
    max_bytes = 2**31 - 1
    try:
        input_size = os.path.getsize(filepath)
        bytes_in = bytearray(0)
        with open(filepath, 'rb') as f_in:
            for _ in range(0, input_size, max_bytes):
                bytes_in += f_in.read(max_bytes)
        obj = cPickle.loads(bytes_in)
    except:
        return None
    return obj
%time df = try_to_load_as_pickled_object_or_None('talkingdata_train.csv.pkl')
CPU times: user 25 s, sys: 1min 13s, total: 1min 38s
Wall time: 2min 21s