python - pyspark 根据两列中的数据创建字典

Question

我有一个包含两列的 pyspark 数据框:

[Row(zip_code='58542', dma='MIN'),
 Row(zip_code='58701', dma='MIN'),
 Row(zip_code='57632', dma='MIN'),
 Row(zip_code='58734', dma='MIN')]

如何从列中的数据中生成键值对？

例如:

{
 "58542":"MIN",
 "58701:"MIN",
 etc..
}

出于性能原因，我想避免使用 collect。我已经尝试了一些方法，但似乎无法只获得值。

Answer 1

您可以 avoid using a udf这里使用 pyspark.sql.functions.struct和 pyspark.sql.functions.to_json (Spark 2.1 及以上版本):

import pyspark.sql.functions as f
from pyspark.sql import Row

data = [
    Row(zip_code='58542', dma='MIN'),
    Row(zip_code='58701', dma='MIN'),
    Row(zip_code='57632', dma='MIN'),
    Row(zip_code='58734', dma='MIN')
]

df = spark.createDataFrame(data)

df.withColumn("json", f.to_json(f.struct("dma", "zip_code"))).show(truncate=False)
#+---+--------+--------------------------------+
#|dma|zip_code|json                            |
#+---+--------+--------------------------------+
#|MIN|58542   |{"dma":"MIN","zip_code":"58542"}|
#|MIN|58701   |{"dma":"MIN","zip_code":"58701"}|
#|MIN|57632   |{"dma":"MIN","zip_code":"57632"}|
#|MIN|58734   |{"dma":"MIN","zip_code":"58734"}|
#+---+--------+--------------------------------+

如果您希望 zip_code 成为键，您可以直接使用 pyspark.sql.functions.create_map 创建一个 MapType :

df.withColumn("json", f.create_map(["zip_code", "dma"])).show(truncate=False)
#+---+--------+-----------------+
#|dma|zip_code|json             |
#+---+--------+-----------------+
#|MIN|58542   |Map(58542 -> MIN)|
#|MIN|58701   |Map(58701 -> MIN)|
#|MIN|57632   |Map(57632 -> MIN)|
#|MIN|58734   |Map(58734 -> MIN)|
#+---+--------+-----------------+