在PySpark中加載存儲在npz存檔中的numpy數組

Question

我在npz存檔中存儲了大量S3中的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)