如何将巨大的 2D NumPy 数组写入缓冲区

Posted 2023-02-16

【中文标题】如何将巨大的 2D NumPy 数组写入缓冲区

【英文标题】：How to write a huge 2D NumPy array into a buffer

【发布时间】：2021-02-07 12:05:12

【问题描述】：

我有一个巨大的 2D numpy 数组 (dtype=bool) 和一个缓冲区，我想将这个 2D 数组写入缓冲区。目前，我执行以下操作，

# Python version 3.7.7, NumPy version 1.18.5
# shape in the dummy_array is just an example, sometimes will be bigger
dummy_array = np.array(np.empty((599066148, 213), dtype='bool'), dtype='bool')

# Pyarrow plasma store buffer
buf = client.create(object_id, dummy_array.nbytes)

# Get a NumPy view of the buffer
array = np.frombuffer(buf, dtype="bool").reshape(dummy_array.shape)

# Write the data or the NumPy array to the buffer
array[:] = dummy_array

问题是这至少需要 3 分钟。 dummy_array 的大小通常为 100 到 200GB，有时甚至更多。我无法弄清楚如何使用memoryview 和np.ctypeslib.as_array(buf, shape=dummy_array.shape) 来执行此操作，正如question 中提到的二维数组（我试过，但没有用）。任何以更好或更快的方式执行此操作的指针都会很棒，因为我将这样做至少数百次，因此，每次迭代节省 30 到 60 秒甚至可以节省数小时。

【问题讨论】：

那么您是在具有 200+ GB RAM 的节点上运行它的吗？

@MateenUlhaq 是的，一个计算服务器。它有 1TB 内存。

【参考方案1】：

您不能分配多维内存视图切片。

NotImplementedError: memoryview slice assignments are currently restricted to ndim = 1

因此，在将数组复制到内存视图之前，您可能需要将数组重塑为一维。

>>> dummy_array = np.array(np.empty((2, 213), dtype='bool'), dtype='bool').reshape(2*213)
>>> mem = memoryview(dummy_array)
>>> mem[0]
True
>>> np.frombuffer(mem, dtype="bool").reshape(dummy_array.shape)
array([ True,  True,  True,  True,  True,  True, False, False,  True,

如果你尝试使用多维，你会得到这个错误。

>>> dummy_array = np.array(np.empty((2, 213), dtype='bool'), dtype='bool')
>>> mem = memoryview(dummy_array)
>>> mem[0]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
NotImplementedError: multi-dimensional sub-views are not implemented

我无法告诉您这是否会比您的其他方法更快，但它可能会为您提供一些关于如何使 memoryview 版本正常工作的想法。