Transform RDD в PySpark

Question

Например, у меня есть следующий RDD типа ((i,j), k):

((0,0), 0) 
((0,1), 0) 
((1,0), 0) 
((1,1), 0) 

Я хочу, чтобы преобразовать его в другой, который имеет 1 если i==j. Моя первая попытка идет не так:

rddnew = rdd.flatMap(lambda ((i,j), k): [if i==j: ((i,j), 1)])) 

Может ли кто-нибудь помочь мне улучшить этот код в python?

Answer 1

Вот решение:

data = [((0, 0), 0), ((0, 1), 0), ((1, 0), 0), ((1, 1), 0)] 
rdd = sc.parallelize(data) 
rdd2 = rdd.map(lambda ((i, j), _): ((i, j), 1) if (i == j) else ((i, j), 0)) 
rdd2.collect() 
# [((0, 0), 1), ((0, 1), 0), ((1, 0), 0), ((1, 1), 1)] 

Вы можете также определить более чистое решение, используя функцию на картографа:

def transformation(entry): 
    (i, j), v = entry 
    return (i, j), v + 1 if i == j else 0 

rdd3 = rdd.map(transformation) 
rdd3.collect() 
# [((0, 0), 1), ((0, 1), 0), ((1, 0), 0), ((1, 1), 1)]