【问题标题】：pyspark create dictionary from data in two columnspyspark 从两列中的数据创建字典
【发布时间】：2019-02-09 21:14:57
【问题描述】：

我有一个包含两列的 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。我尝试了一些方法，但似乎无法仅获得值。

【问题讨论】：

标签： python pyspark

【解决方案1】：

我不确定这与其他解决方案相比性能如何，但您可以这样做：

dict = {row['zipcode']:row['dma'] for row in df.collect()}
print(dict)
#{'58542': 'MIN', '58701': 'MIN', '57632': 'MIN', '58734': 'MIN'}

【讨论】：

【解决方案2】：

还有另一种方法可以将数据框转换为 dict。为此，您需要将数据帧转换为键值对 rdd，因为它仅适用于键值对 rdd。因为字典本身就是键值对的组合。

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')
]

>>> data.show();
+---+--------+
|dma|zip_code|
+---+--------+
|MIN|   58542|
|MIN|   58701|
|MIN|   57632|
|MIN|   58734|
+---+--------+

将您的数据框转换为 rdd。

newrdd = data.rdd

因为您希望 zip_code 作为键，dma 作为值，所以选择 rdd 元素“1”作为键，元素“0”作为值。

keypair_rdd = newrdd.map(lambda x : (x[1],x[0]))

一旦您拥有密钥对 rdd，然后只需使用 collectAsMap 将其转换为字典

>>> dict = keypair_rdd.collectAsMap()
>>> print dict
{u'58542': u'MIN', u'57632': u'MIN', u'58734': u'MIN', u'58701': u'MIN'}

>>> dict.keys()
[u'58542', u'57632', u'58734', u'58701']

查找特定键的值：

>>> dict.get('58542')
u'MIN'

【讨论】：

【解决方案3】：

您可以在此处使用pyspark.sql.functions.struct 和pyspark.sql.functions.to_json avoid using a udf（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)|
#+---+--------+-----------------+

【讨论】：

【解决方案4】：

正如 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"}}']

【讨论】：

创建地图不需要udf。

【解决方案5】：

是的，你可以使用

pyspark.sql.types.MapType(keyType, valueType, valueContainsNull=True)

请分享更多信息，例如数据帧示例输出以及您想要的输出方式，这将有助于编写相同的代码 sn-p。

【讨论】：