带有 json 的 Pyspark 数据框，迭代以创建新的数据框

Posted 2023-04-15

【中文标题】带有 json 的 Pyspark 数据框，迭代以创建新的数据框

【英文标题】：Pyspark dataframe with json, iteration to create new dataframe

【发布时间】：2021-04-28 19:13:26

【问题描述】：

我有以下格式的数据：

customer_id	model
1	[color: 'red', group: 'A',color: 'green', group: 'B']
2	[color: 'red', group: 'A']

我需要处理它，以便创建一个具有以下输出的新数据框：

customer_id	color	group
1	red	A
1	green	B
2	red	A

现在我可以用 python 轻松做到这一点：

import pandas as pd
import json

newdf = pd.DataFrame([])

for index, row in df.iterrows():
    s = row['model']
    x = json.loads(s)
    
    colors_list = []
    users_list = []
    groups_list = []
    
    for i in range(len(x)):
        colors_list.append(x[i]['color'])
        users_list.append(row['user_id'])
        groups_list.append(x[i]['group'])
        
    newdf = newdf.append(pd.DataFrame('customer_id': users_list, 'group': groups_list, 'color': colors_list))

如何使用 pyspark 获得相同的结果？

我正在显示原始数据框的第一行和架构：

+-----------+--------------------+
|customer_id|              model |
+-----------+--------------------+
|       3541|["score":0.04767...|
|     171811|["score":0.04473...|
|      12008|["score":0.08043...|
|      78964|["score":0.06669...|
|     119600|["score":0.06703...|
+-----------+--------------------+
only showing top 5 rows

root
 |-- user_id: integer (nullable = true)
 |-- groups: string (nullable = true)

【参考方案1】：

from_json可以解析出包含Json数据的字符串列：

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

data = [[1, "[color: 'red', group: 'A',color: 'green', group: 'B']"],
        [2, "[color: 'red', group: 'A']"]]

df = spark.createDataFrame(data, schema = ["customer_id", "model"]) \
    .withColumn("model", F.from_json("model", T.ArrayType(T.MapType(T.StringType(), T.StringType())), "allowUnquotedFieldNames": True)) \
    .withColumn("model", F.explode("model")) \
    .withColumn("color", F.col("model")["color"]) \
    .withColumn("group", F.col("model")["group"]) \
    .drop("model")

结果：

+-----------+-----+-----+
|customer_id|color|group|
+-----------+-----+-----+
|          1|  red|    A|
|          1|green|    B|
|          2|  red|    A|
+-----------+-----+-----+

【讨论】：

谢谢。如何实现数据中的列表列表？数据框有 500 万行，当我尝试以该格式连接两列然后按照您的建议对其进行处理时，出现错误：无法推断类型的模式：

@Sapehi 我的回答假定输入数据集的 model 字段中的所有 json 字符串都具有相同的模式：颜色/组组合的数组。如果该列中还有其他 json 字符串，也许您可​​以发布一个示例？

我认为这是因为我在 spark 数据框中拥有原始数据，并且我不知道如何将这两列作为列表列表，就像您在变量“数据”中一样。当我尝试时，我得到 [Row(customer_id=7286, groups=u'["color":'red'.....]，这种格式稍后会给我错误。

@Sapehi 代码应该可以在 Spark 数据帧中正常工作。您能否在问题中包含originalData.show() 和originalData.printSchema() 的输出？也许这个

谢谢。我已添加该信息以查看是否有帮助:)