Как выровнять элементы в списке списков, используя Pyspark?

Question

I/p: l=[[1,2,3],[3,4]]

O/p: [[1,4,9],[9,16]]

Как я могу достичь выше, используя PySpark?

Я пытался

sc.parallelize(l).flatmap(lambda x:x).map(lambda x: x*x).collect ()

И получил ниже вывод

[1,4,9,9,16]

Но мне нужен вывод, какэто

O/p: [[1,4,9],[9,16]]

Answer 1

вы используете flatMap и теряете информацию о каждом наборе, вы должны перебирать каждый список на карте

>>> rdd.map(lambda x: [i * i for i in x]).collect()
[[1, 4, 9], [9, 16]]

Answer 2

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

def square_it(num):
    return [i*i for i in num]

squares = rdd.map(square_it)
result = squares.collect()
print(result)