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 。我已经尝试了一些东西，但似乎无法获得 values 。

原文由 too_many_questions 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

正如 Ankin 所说，您可以为此使用 MapType ：

 import pyspark
from pyspark.sql import Row

sc = pyspark.SparkContext()
spark = pyspark.sql.SparkSession(sc)

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

data.show()

输出：

 +---+--------+
|dma|zip_code|
+---+--------+
|MIN|   58542|
|MIN|   58701|
|MIN|   57632|
|MIN|   58734|
+---+--------+

from pyspark.sql.functions import udf
from pyspark.sql import types as T

@udf(T.MapType(T.StringType(), T.StringType()))
def create_struct(zip_code, dma):
    return {zip_code: dma}

data.withColumn('struct', create_struct(data.zip_code, data.dma)).toJSON().collect()

输出：

 ['{"dma":"MIN","zip_code":"58542","struct":{"58542":"MIN"}}',
 '{"dma":"MIN","zip_code":"58701","struct":{"58701":"MIN"}}',
 '{"dma":"MIN","zip_code":"57632","struct":{"57632":"MIN"}}',
 '{"dma":"MIN","zip_code":"58734","struct":{"58734":"MIN"}}']

原文由 martinarroyo 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

您 可以避免使用 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 成为关键，您可以创建一个 MapType 直接使用 pyspark.sql.functions.create_map ：

 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)|
#+---+--------+-----------------+

原文由 pault 发布，翻译遵循 CC BY-SA 4.0 许可协议