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

Question

我在S3中存储了npz存档中的大量numpy数组。将它们加载到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)