在 Pyspark 中拆分 RDD 分區中的數組

Question

我有一個 3D 數值數據文件，我分塊讀取（因為分塊讀取比單個索引快）。 例如說在“文件”中有一個 MxNx30 數組，我會創建一個這樣的 RDD：

def read(ind):
    f = customFileOpener(file)
    return f['data'][:,:,ind[0]:ind[-1]+1]

indices = [[0,9],[10,19],[20,29]]
rdd = sc.parallelize(indices,3).map(lambda v:read(v))
rdd.count()

所以這 3 個分區中的每一個都有一個大小為 MxNx10 的 numpy.ndarray 元素。

現在，我想拆分這些元素中的每一個，因此在每個分區中，我有 10 個元素，每個元素都是一個 MxN 數組。 為此，我嘗試使用 flatMap()，但收到“NoneType 對象不可迭代”的錯誤：

def splitArr(arr):
    Nmid = arr.shape[-1]
    out = []
    for i in range(0,Nmid):
         out.append(arr[...,i])
    return out

rdd2 = rdd.flatMap(lambda v: splitArr(v))
rdd2.count()

這樣做的正確方法是什么？ 關鍵點是 (a) 我需要從文件中以塊的形式讀取數據和 (b) 拆分數據，以便元素大小為 MxN（最好保持分區結構）。

Answer 1

據我了解，您的描述應該是這樣的：

rdd.flatMap(lambda arr: (x for x in np.rollaxis(arr, 2)))

或者，如果您更喜歡單獨的功能：

def splitArr(arr):
    for x in np.rollaxis(arr, 2):
        yield x

rdd.flatMap(splitArr)