如何从一个元组可迭代对象填充两个（或更多）numpy数组？

Question

我现在遇到的问题是，我想在内存中存储一个很长的有序列表，里面包含的是(浮点数, 字符串)的元组。因为我的内存只有4GB，普通的列表放不下，所以我想用两个numpy.ndarray来解决。

这些数据的来源是一个包含2个元素的元组的可迭代对象。numpy有一个叫fromiter的函数，但我不知道该怎么用。因为我不知道可迭代对象里面有多少个元素，所以不能先把它变成列表，这样会占用太多内存。我想到了itertools.tee，但它似乎会增加很多内存开销。

我想我可以分块处理这个迭代器，把每一块的数据添加到数组里。那么我的问题是，怎么才能高效地做到这一点呢？我是不是应该创建两个二维数组，然后往里面添加行？（之后我还需要把它们转换成一维的）。

或者说有没有更好的方法？我真正需要的就是通过对应的数字值在字符串数组中进行搜索，最好能在对数时间内完成（这就是我想按浮点数值排序的原因），而且还希望能尽量节省内存。

附注：这个可迭代对象是没有排序的。

Answer 1

这里有一种方法，可以从一个生成 N 元组的生成器中创建 N 个独立的数组：

import numpy as np
import itertools as IT


def gendata():
    # You, of course, have a different gendata...
    N = 100
    for i in xrange(N):
        yield (np.random.random(), str(i))


def fromiter(iterable, dtype, chunksize=7):
    chunk = np.fromiter(IT.islice(iterable, chunksize), dtype=dtype)
    result = [chunk[name].copy() for name in chunk.dtype.names]
    size = len(chunk)
    while True:
        chunk = np.fromiter(IT.islice(iterable, chunksize), dtype=dtype)
        N = len(chunk)
        if N == 0:
            break
        newsize = size + N
        for arr, name in zip(result, chunk.dtype.names):
            col = chunk[name]
            arr.resize(newsize, refcheck=0)
            arr[size:] = col
        size = newsize
    return result

x, y = fromiter(gendata(), '<f8,|S20')

order = np.argsort(x)
x = x[order]
y = y[order]

# Some pseudo-random value in x
N = 10
val = x[N]
print(x[N], y[N])
# (0.049875262239617246, '46')

idx = x.searchsorted(val)
print(x[idx], y[idx])
# (0.049875262239617246, '46')

上面的 fromiter 函数会分块读取可迭代对象（每块的大小由 chunksize 决定）。它会调用 NumPy 数组的方法 resize，根据需要扩展结果数组的大小。

我使用了一个较小的默认 chunksize，因为我是在小数据上测试这段代码。当然，你可以选择更改默认的 chunksize，或者传入一个更大的 chunksize 参数。

Answer 2

也许可以使用 np.fromiter 来构建一个结构化的数组：

import numpy as np


def gendata():
    # You, of course, have a different gendata...
    for i in xrange(N):
        yield (np.random.random(), str(i))

N = 100

arr = np.fromiter(gendata(), dtype='<f8,|S20')

通过第一列进行排序，第二列作为平局时的决定因素，这个过程需要 O(N log N) 的时间：

arr.sort(order=['f0','f1'])

根据第一列的值找到对应的行，可以使用 searchsorted，这个过程只需要 O(log N) 的时间：

# Some pseudo-random value in arr['f0']
val = arr['f0'][10]
print(arr[10])
# (0.049875262239617246, '46')

idx = arr['f0'].searchsorted(val)
print(arr[idx])
# (0.049875262239617246, '46')

---

你在评论中问了很多重要的问题；我来这里尝试回答一下：

• 基本的数据类型在 numpybook 中有解释。可能还有一两个额外的数据类型（比如 float16，这是在那本书写完后新增的，但基本概念都在那儿解释了。）

  也许更详细的讨论可以在 在线文档 中找到。这是对你提到的例子的一个很好的补充 这里。

• 数据类型可以用来定义带有列名的结构化数组，或者使用默认的列名。'f0'、'f1' 等是默认的列名。因为我定义的数据类型是 '<f8,|S20'，所以没有提供列名，因此 NumPy 将第一列命名为 'f0'，第二列命名为 'f1'。如果我们使用

  dtype='[('fval','<f8'), ('text','|S20')]
  

  那么结构化数组 arr 的列名将会是 'fval' 和 'text'。

• 不幸的是，数据类型必须在调用 np.fromiter 时就固定下来。你可以先遍历一次 gendata 来找出字符串的最大长度，构建你的数据类型，然后再调用 np.fromiter（再遍历一次 gendata），但这样做比较麻烦。当然，如果你提前知道字符串的最大长度会更好。（|S20 定义了字符串字段的固定长度为 20 字节。）

• NumPy 数组将预定义大小的数据放入固定大小的数组中。可以把数组（即使是多维的）想象成一个连续的一维内存块。（这有点过于简化——实际上有非连续的数组——但这有助于你理解接下来的内容。）NumPy 的速度很大程度上是通过利用固定大小（由 dtype 设置）来快速计算访问数组元素所需的偏移量。如果字符串的大小是可变的，那么 NumPy 就很难找到正确的偏移量。这里的“难”是指 NumPy 需要一个索引，或者需要重新设计。NumPy 本身并不是这样构建的。

• NumPy 确实有一个 object 数据类型，这允许你放置一个 4 字节的指针指向任何你想要的 Python 对象。这样，你就可以有包含任意 Python 数据的 NumPy 数组。不幸的是，np.fromiter 函数不允许你创建 object 类型的数组。我不太确定为什么会有这个限制……

• 注意，当指定 count 时，np.fromiter 的性能会更好。通过知道 count（行数）和 dtype（因此每行的大小），NumPy 可以预先分配足够的内存来存放结果数组。如果你不指定 count，NumPy 会猜测数组的初始大小，如果太小，它会尝试调整数组的大小。如果原来的内存块可以扩展，那你就幸运了。但如果 NumPy 必须分配一个全新的内存块，那么所有旧数据都必须复制到新位置，这会显著降低性能。