如何用不相等的列表平面映射数据框?

时间:2018-07-02 13:57:41

标签: python pyspark

我的数据框如下-

  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 

列数不是固定的。如果该解决方案是广义的,则将有所帮助。

1 个答案:

答案 0 :(得分:0)

一种选择是使用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列中的数组填充为适当的长度。