Python
Python3

Python3 & MacOSで4GB以上のオブジェクトをpickle化するときのエラーと対処法

タイトルそのままですが、Python3でMacOSで4GB以上のオブジェクトをpickle化すると

import pickle

path = "path.pkl"
data_larger_than_4GB = bytearray(2**32)
pickle.dump(data_larger_than_4GB, open(path, 'wb'))

次のようなエラーが出ます。

OverflowError: cannot serialize a bytes object larger than 4 GiB

機械学習の学習データなど4GBとかすぐ超えるのでこれでは使い物になりません。これに対して、

pickle.dump(data_larger_than_4GB, open(path, 'wb'), protocol=4)

protocol=4という引数をしてしてやると解決するというStackoverflowの回答があります。これはWindowsであればこれで解決するみたいですが、Macの場合はMac特有のバグがあるらしくこれでもエラーがでます。

OSError: [Errno 22] Invalid argument

これ問題に対する完璧な解決がコード付きで別のStackoverflowに解決があります。

class MacOSFile(object):

    def __init__(self, f):
        self.f = f

    def __getattr__(self, item):
        return getattr(self.f, item)

    def read(self, n):
        # print("reading total_bytes=%s" % n, flush=True)
        if n >= (1 << 31):
            buffer = bytearray(n)
            idx = 0
            while idx < n:
                batch_size = min(n - idx, 1 << 31 - 1)
                # print("reading bytes [%s,%s)..." % (idx, idx + batch_size), end="", flush=True)
                buffer[idx:idx + batch_size] = self.f.read(batch_size)
                # print("done.", flush=True)
                idx += batch_size
            return buffer
        return self.f.read(n)

    def write(self, buffer):
        n = len(buffer)
        print("writing total_bytes=%s..." % n, flush=True)
        idx = 0
        while idx < n:
            batch_size = min(n - idx, 1 << 31 - 1)
            print("writing bytes [%s, %s)... " % (idx, idx + batch_size), end="", flush=True)
            self.f.write(buffer[idx:idx + batch_size])
            print("done.", flush=True)
            idx += batch_size


def pickle_dump(obj, file_path):
    with open(file_path, "wb") as f:
        return pickle.dump(obj, MacOSFile(f), protocol=pickle.HIGHEST_PROTOCOL)


def pickle_load(file_path):
    with open(file_path, "rb") as f:
        return pickle.load(MacOSFile(f))

pickle.dump pickle.loadの代わりにpickle_dump pickle_loadを使ってやれば、バグを回避できます。すばらしい。