python - 将 numpy 数组从 GCS 读取到 Spark

Question

我在谷歌存储中有 100 个包含 numpy 数组的 npz 文件。我已经使用 jupyter 设置了 dataproc，并且正在尝试将所有 numpy 数组读入 Spark RDD。将 numpy 数组从 google 存储加载到 pyspark 的最佳方法是什么？有没有像 np.load("gs://path/to/array.npz") 这样的简单方法加载 numpy 数组，然后执行 sc.parallelize就在上面吗？

Answer 1

如果您计划最终扩展，您将需要使用SparkContext中的分布式输入方法。而不是依赖 sc.parallelize 从驱动程序加载任何本地文件。听起来您需要完整地读取每个文件，所以在您的情况下您想要:

npz_rdd = sc.binaryFiles("gs://path/to/directory/containing/files/")

或者，如果需要，您也可以指定单个文件，但这样您就只有一个包含单个元素的 RDD:

npz_rdd = sc.binaryFiles("gs://path/to/directory/containing/files/arr1.npz")

那么每条记录都是一对<filename>,<str of bytes> 。在 Dataproc 上，sc.binaryFiles将自动直接使用 GCS 路径，与 np.load 不同这需要本地文件系统路径。

然后在你的工作代码中，你只需要使用 StringIO使用这些字节字符串作为您放入 np.load 的文件对象:

from StringIO import StringIO
# For example, to create an RDD of the 'arr_0' element of each of the picked objects:
npz_rdd.map(lambda l: numpy.load(StringIO(l[1]))['arr_0'])

在开发过程中，如果您确实只想将文件读入主驱动程序，您可以随时使用collect()折叠您的RDD。在本地检索它:

npz_rdd = sc.binaryFiles("gs://path/to/directory/containing/files/arr1.npz")
local_bytes = npz_rdd.collect()[0][1]
local_np_obj = np.load(StringIO(local_bytes))