使用pyspark，如何在保留其他列的同时将包含变量映射的列扩展到DataFrame中的新列？

Question

我有以下数据帧

+-----+--------------------------------------------------+---+
|asset|signals                                           |ts |
+-----+--------------------------------------------------+---+
|2    |[D -> 1100, F -> 3000]                            |6  |
|1    |[D -> 500, System.Date -> 340]                    |5  |
|1    |[B -> 100, E -> 900, System.Date -> 310]          |4  |
|1    |[B -> 110, C -> 200, System.Date -> 320]          |3  |
|1    |[A -> 330, B -> 120, C -> 210, D -> 410, E -> 100]|2  |
+-----+--------------------------------------------------+---+

我需要将具有键值的 column:'signals' 投影到多个列，如下所示：

+-----+---+-----------+----+----+----+----+----+----+
|asset|ts |System.Date|F   |E   |B   |D   |C   |A   |
+-----+---+-----------+----+----+----+----+----+----+
|2    |6  |null       |3000|null|null|1100|null|null|
|1    |5  |340        |null|null|null|500 |null|null|
|1    |4  |310        |null|900 |100 |null|null|null|
|1    |3  |320        |null|null|110 |null|200 |null|
|1    |2  |null       |null|100 |120 |410 |210 |330 |
+-----+---+-----------+----+----+----+----+----+----+

所以这里是例子：

d = [{'asset': '2', 'ts': 6, 'signals':{'F': '3000','D':'1100'}}, 
     {'asset': '1', 'ts': 5, 'signals':{'System.Date': '340','D':'500'}}, 
     {'asset': '1', 'ts': 4, 'signals':{'System.Date': '310', 'B': '100', 'E':'900'}}, 
     {'asset': '1', 'ts': 3, 'signals':{'System.Date': '320', 'B': '110','C':'200'}},
      {'asset': '1', 'ts': 2, 'signals':{'A': '330', 'B': '120','C':'210','D':'410','E':'100'}}]
df = spark.createDataFrame(d)

我可以提取所有可能的密钥并实现我的目标，如下所示：

from pyspark.sql import functions as spfn
# the following takes too long (WANT-TO-AVOID)
all_signals = (df
    .select(spfn.explode("signals"))
    .select("key")
    .distinct()
    .rdd.flatMap(lambda x: x)
    .collect())
print(all_signals)
exprs = [spfn.col("signals").getItem(k).alias(k) for k in all_signals]

df1 = df.select(spfn.col('*'),*exprs).drop('signals')

df1.show(truncate=False)

['System.Date', 'F', 'E', 'B', 'D', 'C', 'A']
+-----+---+-----------+----+----+----+----+----+----+
|asset|ts |System.Date|F   |E   |B   |D   |C   |A   |
+-----+---+-----------+----+----+----+----+----+----+
|2    |6  |null       |3000|null|null|1100|null|null|
|1    |5  |340        |null|null|null|500 |null|null|
|1    |4  |310        |null|900 |100 |null|null|null|
|1    |3  |320        |null|null|110 |null|200 |null|
|1    |2  |null       |null|100 |120 |410 |210 |330 |
+-----+---+-----------+----+----+----+----+----+----+

但我想知道是否有办法使用以下内容但不知道如何保留现有列：

df2 = spark.read.json(df.rdd.map(lambda r: r.signals))
df2.show(truncate=False)
+----+----+----+----+----+----+-----------+
|A   |B   |C   |D   |E   |F   |System.Date|
+----+----+----+----+----+----+-----------+
|null|null|null|1100|null|1000|null       |
|null|null|null|500 |null|null|340        |
|null|100 |null|null|900 |null|310        |
|null|110 |200 |null|null|null|320        |
|330 |120 |210 |410 |100 |null|null       |
+----+----+----+----+----+----+-----------+

获取所有密钥的上述步骤（标记为 WANT-TO-AVOID）需要很长时间，而“spark.read.json”似乎要快得多。那么，有没有更简单快捷的方法来扩展地图列？

Answer 1

您可以执行此操作read.json以获得所需的列。(spark2.4+)。

from pyspark.sql import functions as F

df1=df.withColumn("signals", F.map_concat("signals", F.create_map(F.lit("asset"),"asset",F.lit("ts"),"ts")))

df2 = spark.read.json(df1.rdd.map(lambda r: r.signals))

df2.show()

#+----+----+----+----+----+----+-----------+-----+---+
#|   A|   B|   C|   D|   E|   F|System.Date|asset| ts|
#+----+----+----+----+----+----+-----------+-----+---+
#|null|null|null|1100|null|3000|       null|    2|  6|
#|null|null|null| 500|null|null|        340|    1|  5|
#|null| 100|null|null| 900|null|        310|    1|  4|
#|null| 110| 200|null|null|null|        320|    1|  3|
#| 330| 120| 210| 410| 100|null|       null|    1|  2|
#+----+----+----+----+----+----+-----------+-----+---+