python - Unpickling 中流(python)

Question

我正在编写脚本来通过反复取消对象直到 EOF 来处理(非常大的)文件。我想对文件进行分区并让单独的进程(在云中)解开并处理单独的部分。

但是我的分区器并不智能，它不知道文件中 pickle 对象之间的边界(因为这些边界取决于被 pickle 的对象类型等)。

有没有办法扫描文件中的“开始 pickle 对象”哨兵？天真的方法是尝试在连续的字节偏移处进行 unpickling，直到对象被成功 pickle 为止，但这会产生意想不到的错误。似乎对于某些输入组合，unpickler 不同步并且不会为文件的其余部分返回任何内容(请参见下面的代码)。

import cPickle
import os

def stream_unpickle(file_obj):
    while True:
        start_pos = file_obj.tell()
        try:
            yield cPickle.load(file_obj)
        except (EOFError, KeyboardInterrupt):
            break
        except (cPickle.UnpicklingError, ValueError, KeyError, TypeError, ImportError):
            file_obj.seek(start_pos+1, os.SEEK_SET)

if __name__ == '__main__':
    import random
    from StringIO import StringIO

    # create some data
    sio = StringIO()
    [cPickle.dump(random.random(), sio, cPickle.HIGHEST_PROTOCOL) for _ in xrange(1000)]
    sio.flush()

    # read from subsequent offsets and find discontinuous jumps in object count
    size = sio.tell()
    last_count = None
    for step in xrange(size):
        sio.seek(step, os.SEEK_SET)
        count = sum(1 for _ in stream_unpickle(file_obj))
        if last_count is None or count == last_count - 1:
            last_count = count
        elif count != last_count:
            # if successful, these should never print (but they do...)
            print '%d elements read from byte %d' % (count, step)
            print '(%d elements read from byte %d)' % (last_count, step-1)
            last_count = count

Answer 1

pickletools 模块有一个显示操作码的 dis 函数。它表明有一个您可能正在扫描的 STOP 操作码:

>>> import pickle, pickletools, StringIO
>>> s = StringIO.StringIO()
>>> pickle.dump('abc', s)
>>> p = s.getvalue()
>>> pickletools.dis(p)
    0: S    STRING     'abc'
    7: p    PUT        0
   10: .    STOP
highest protocol among opcodes = 0

请注意，使用 STOP 操作码有点棘手，因为代码的长度可变，但它可以作为有关截止位置的有用提示。

如果您控制另一端的 pickle 步骤，则可以通过添加您自己明确的替代分隔符来改善这种情况:

>>> sep = '\xDE\xAD\xBE\xEF'
>>> s = StringIO.StringIO()
>>> pickle.dump('abc', s)
>>> s.write(sep)
>>> pickle.dump([10, 20], s)
>>> s.write(sep)
>>> pickle.dump('def', s)
>>> s.write(sep)
>>> pickle.dump([30, 40], s)
>>> p = s.getvalue()

在拆包之前，使用已知的分隔符分成单独的 pickle :

>>> for pick in p.split(sep):
        print pickle.loads(pick)

abc
[10, 20]
def
[30, 40]