在 PySpark 中加载存储在 npz 存档中的 numpy 数组

Question

我在 s3 中有大量 numpy 数组存储在 npz 存档中。将它们加载到 PySpark RDD/NumPy 数组的数据帧中的最佳方法是什么？我尝试使用 sc.wholeTextFiles API 加载文件。

rdd=sc.wholeTextFiles("s3://[bucket]/[folder_containing_npz_files]") 

但是numpy.load需要一个文件句柄。将文件内容作为字符串加载到内存中会占用大量内存。

Answer 1

你不能对内存要求做太多事情，但除此之外BytesIO应该可以正常工作：

from io import BytesIO
def extract(kv):
    k, v = kv
    with BytesIO(v) as r:
        for f, x in np.load(r).items():
            yield "{0}t{1}".format(k, f), x
sc.binaryFiles(inputPath).flatMap(extract)