Как выровнять карту данных с неравными списками?

Question

У меня есть кадр данных следующим образом -

  a  |       b       |   c      
 [1] |    [3,4,5,6] | [7,8,9,10]

мне нужен вывод как

  a  |   b       |   c
  1      3           7
  1      4           8
  1      5           9
  1      6           10 

В настоящее время я получаю вывод, как показано ниже, используя следующее утверждение -

cols=['a','b','c']
df.rdd.flatMap(lambda x: itertools.izip_longest(*[x[c] for c in cols])).toDF(cols)

  a  |   b       |   c
  1      3           7
 null    4           8
 null    5           9
 null    6           10 

Количество столбцов не фиксировано. Если бы решение было обобщенным, это было бы полезно.

Answer 1

Один из вариантов - использовать itertools.repeat для повторения каждого массива, длина которого меньше длины самого длинного массива.

from itertools import izip_longest, repeat, chain

cols = df.columns
def zip_with_repeat(row, cols):
    M = max(len(row[c]) for c in cols)
    return izip_longest(
        *[list(chain(*repeat(row[c], (M - len(row[c])) + 1)))[:M] for c in cols]
    )

df.rdd.flatMap(lambda row: zip_with_repeat(row, cols)).toDF(cols).show()
#+---+---+---+
#|  a|  b|  c|
#+---+---+---+
#|  1|  3|  7|
#|  1|  4|  8|
#|  1|  5|  9|
#|  1|  6| 10|
#+---+---+---+

Для наглядности предположим, что вместо этого у вас был следующий DataFrame:

#+--------+------------+-------------+
#|       a|           b|            c|
#+--------+------------+-------------+
#|     [1]|[3, 4, 5, 6]|[7, 8, 9, 10]|
#|[10, 20]|[30, 40, 50]| [70, 80, 90]|
#+--------+------------+-------------+

Код выдаст:

#+---+---+---+
#|  a|  b|  c|
#+---+---+---+
#|  1|  3|  7|
#|  1|  4|  8|
#|  1|  5|  9|
#|  1|  6| 10|
#| 10| 30| 70|
#| 20| 40| 80|
#| 10| 50| 90|
#+---+---+---+

Обратите внимание, что 10 повторяется один раз, чтобы заполнить массив в столбце a до соответствующей длины.