高效地将numpy数组写入二进制文件

Question

我需要一种高效的解决方案来将大量数据写入二进制文件。目前我使用numpy的.tofile方法，这占用了大部分运行时间。以下是我的MWE：

import numpy as np

def writeCFloat(f, ndarray):
    np.asarray(ndarray, dtype=np.float32).tofile(f)
def writeCInt(f, ndarray):
    np.asarray(ndarray, dtype=np.int32).tofile(f)
def writeC80(f, string):
    np.asarray(string, dtype='a80').tofile(f)

if __name__ == "__main__":
    f = open('test.out', mode='wb')
    ndarray = np.zeros((10000,10000))

    writeCInt(f, ndarray)
    writeCFloat(f, ndarray)
    writeC80(f, 'coordinates')

那么，有比现有解决方案更有效的方法吗？提前感谢。

- Magdalena

---

考虑到您的代码中没有显式循环，以及 asarray 的文档中提到 如果输入已经是 ndarray，则不执行复制，我认为这可能是最快的方法。 - mkrieger1

为什么当前的解决方案太慢？你有展示Python可以实现更高IO速率的基准测试吗？还是这是一个需要实时写入大量数据的脚本？ - user2699

感谢您的回复。目前我正在读取一个大型二进制文件（约40 GB），然后将数据写回另一个二进制文件。我对Python脚本进行了分析，发现大部分时间都花在了“.tofile”上。因此，我想知道是否有其他方法可以将数组、字符串等写入二进制文件以提高效率？ - Magdalena

也许 dask 对你来说是一个选择？通常并行化可以提高速度，问题在于如何实现它，但这取决于你使用情况的细节。 - a_guest

看起来对我来说没问题。您还可以考虑pandas提供的各种选项（比如HDF、feather或parquet）和dask。 - JohnE

考虑并行性，另一个选项是使用MPI的文件IO功能。 - a_guest

Answer 1

您可以使用dask并行运行此操作。这也使您能够扩展到单个线程的内存限制之外。

根据您的用例和要将数据保存为的文件类型，您可以执行以下操作：

MCVE

import numpy as np
import dask.array as da

def writeCFloat_dask(ndarray):
    da.from_array(ndarray).astype('float32').to_zarr('float32.zarr')
def writeCInt_dask(ndarray):
    da.from_array(ndarray).astype('int32').to_zarr('int32.zarr')
def writeC80_dask(string):
    da.from_array(string).astype('a80').to_zarr('a80.zarr')


if __name__ == "__main__":
    ndarray = np.zeros((100000,10000))

    writeCInt_dask(ndarray)
    writeCFloat_dask(ndarray)
    writeC80_dask('coordinates')

或者，如果您想将所有数组写入单个文件，您可以使用以下方式写入HDF5：

def returnCFloat_dask(ndarray):
    return da.from_array(ndarray).astype('float32')
def returnCInt_dask(ndarray):
    return da.from_array(ndarray).astype('int32')
def returnC80_dask(string):
    return da.from_array(string).astype('a80')

if __name__ == "__main__":
    ndarray = np.zeros((100000,10000))

    da_float = returnCFloat_dask(ndarray)
    da_int = returnCInt_dask(ndarray)
    da_string = returnC80_dask('coordinates')

    da.to_hdf5('test-out.hdf5', {'/float': da_float, '/int': da_int, '/string': da_string})

请参阅Dask编写数组到二进制文件类型的文档以获取更多信息。

- rrpelgrim