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

Loading numpy arrays stored in npz archive in PySpark

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

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

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

您不能对内存要求做太多,但 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)