python之在 Pyspark 的 RDD 分区中拆分数组

我有一个 3D 数值数据文件，我从中读取数据 block (因为读取数据 block 比单个索引更快)。例如，假设"file"中有一个 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(最好保持分区结构)。

请您参考如下方法:

据我了解，您的描述应该是这样的:

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)