如何从pickle文件中逐行加载数据？

Question

我有一个大的数据集：20,000 x 40,000 的 numpy 数组。我已将其保存为 pickle 文件。

我希望不用一次性读入整个数据集到内存中，而是每次只读取其中几行（比如100行）作为小批量使用。

如何从 pickle 文件中随机选择（且不重复选择）几行数据进行读取？

- StatsSorceress

---

将其存储为其他格式，以便进行随机或增量访问。 - martineau

这并不是您问题的答案，但请记住攻击的漏洞。文档中的警告 - pickle模块不安全，可能受到错误或恶意构造数据的影响。永远不要对来自不受信任或未经身份验证的来源的数据进行反序列化处理。 - MartinP

如果它是一个数字数组，你可以将其存储为二进制文件，并使用 file.seek() 访问文件中的任何一行。struct 模块可用于编写和读取文件。 - martineau

谢谢 Padraic，但有没有一种格式不会加载得很慢？我尝试使用 pandas dataframe 和 csv，但它花了很长时间。 - StatsSorceress

或许 memmap 更符合您的要求。您可以查看 http://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.memmap.html 了解更多信息。 - Padraic Cunningham

Answer 1

你可以逐步地向文件中写入pickle数据，这样你也可以逐步地加载它们。

以以下示例为例。在这里，我们迭代一个列表中的项，并依次对每个项进行pickle。

>>> import cPickle
>>> myData = [1, 2, 3]
>>> f = open('mydata.pkl', 'wb')
>>> pickler = cPickle.Pickler(f)
>>> for e in myData:
...     pickler.dump(e)
<cPickle.Pickler object at 0x7f3849818f68>
<cPickle.Pickler object at 0x7f3849818f68>
<cPickle.Pickler object at 0x7f3849818f68>
>>> f.close()

现在我们可以按照相反的过程逐个加载每个对象。以示例为例，假设我们只需要第一个项目，而不想遍历整个文件。

>>> f = open('mydata.pkl', 'rb')
>>> unpickler = cPickle.Unpickler(f)
>>> unpickler.load()
1

此时，文件流仅前进到第一个对象。未加载剩余的对象，这正是您想要的行为。您可以尝试读取文件的其余部分，并查看其余部分是否仍在原地。

>>> f.read()
'I2\n.I3\n.'

- Alex Smith

---

这是一个优雅的解决方案。谢谢你指出来！ - jonathanking

Answer 2

由于您不了解pickle的内部工作方式，因此需要使用另一种存储方法。下面的脚本使用函数将数据逐行保存在原始文件中。

由于每行的长度已知，因此可以计算它在文件中的偏移量，并通过和进行访问。之后，使用函数将其转换回数组。

然而，需要注意的是，数组的大小未保存（这也可以添加，但需要更多的复杂处理），并且此方法可能不像pickle数组那样可移植。

正如@PadraicCunningham在他的comment中指出的那样，memmap很可能是一个替代方案和优雅的解决方案。

关于性能的备注：在阅读了评论后，我进行了一项简短的基准测试。在我的机器上（16GB RAM，加密SSD），我能够在24秒内完成40000个随机行读取操作（当然，使用的是一个20000x40000的矩阵，而不是例子中的10x10矩阵）。

from __future__ import print_function
import numpy
import random

def dumparray(a, path):
    lines, _ = a.shape
    with open(path, 'wb') as fd:
        for i in range(lines):
            fd.write(a[i,...].tobytes())

class RandomLineAccess(object):
    def __init__(self, path, cols, dtype):
        self.dtype = dtype
        self.fd = open(path, 'rb')
        self.line_length = cols*dtype.itemsize

    def read_line(self, line):
        offset = line*self.line_length
        self.fd.seek(offset)
        data = self.fd.read(self.line_length)

        return numpy.frombuffer(data, self.dtype)

    def close(self):
        self.fd.close()


def main():
    lines = 10
    cols = 10
    path = '/tmp/array'

    a = numpy.zeros((lines, cols))
    dtype = a.dtype

    for i in range(lines):
        # add some data to distinguish lines
        numpy.ndarray.fill(a[i,...], i)

    dumparray(a, path)
    rla = RandomLineAccess(path, cols, dtype)

    line_indices = list(range(lines))
    for _ in range(20):
        line_index = random.choice(line_indices)
        print(line_index, rla.read_line(line_index))

if __name__ == '__main__':
    main()

- code_onkel

Answer 3

谢谢大家。最后我找到了一个解决方法（使用内存更大的机器，这样我就能够将数据集加载到内存中了）。

- StatsSorceress