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

Question

我在 S3 中有大量 numpy 数组存储在 npz 存档中。将它们加载到 NumPy 数组的 PySpark RDD/Dataframe 中的最佳方法是什么？我尝试使用 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)