如何从 PySpark 中的多个列创建字典列表，其中键是列名，值是该列的值？

Posted 2023-04-15

【中文标题】如何从 PySpark 中的多个列创建字典列表，其中键是列名，值是该列的值？

【英文标题】：How to create list of dictionaries from multiple columns in PySpark where key is a column name and value is that column's value?

【发布时间】：2021-03-24 19:14:32

【问题描述】：

考虑一个示例数据框，比如df：

user_id | item_id | item_param1 | item_param2 |
  1          a           10            20
  1          b           30            40
  2          b           50            60
  2          c           70            80

从这个数据框df，我想在一行中为每个user_id 收集所有items 及其属性item_paramX，因此输出将是：

user_id |    values
------------------------------------------------
  1     |   ['a': 'item_param1': 10, 'item_param2': 20, 'b': 'item_param1': 30, 'item_param2': 40]
  2     |   ['b': 'item_param1': 50, 'item_param2': 60, 'c': 'item_param1': 70, 'item_param2': 80]

我需要使用groupBy('user_id')，然后是某种形式的collect_list 来获取字典列表。

【参考方案1】：

你可以在item_id -> struct的地图上使用collect_list：

import pyspark.sql.functions as F

df2 = df.groupBy('user_id').agg(
    F.to_json(
        F.collect_list(
            F.create_map(
                'item_id', 
                F.struct('item_param1', 'item_param2')
            )
        )
    ).alias('values')
)

df2.show(truncate=False)
+-------+-------------------------------------------------------------------------------------+
|user_id|values                                                                               |
+-------+-------------------------------------------------------------------------------------+
|1      |["a":"item_param1":10,"item_param2":20,"b":"item_param1":30,"item_param2":40]|
|2      |["b":"item_param1":50,"item_param2":60,"c":"item_param1":70,"item_param2":80]|
+-------+-------------------------------------------------------------------------------------+