PySpark：将字典数据附加到 PySpark DataFrame

Posted 2023-04-15

【中文标题】PySpark：将字典数据附加到 PySpark DataFrame

【英文标题】：PySpark: Attach dictionary data to PySpark DataFrame

【发布时间】：2018-06-19 10:51:51

【问题描述】：

我的问题的简化版本是这样的：

我有一个 Spark DataFrame ("my_df")，它有一列 ("col1") 和值 'a','b','c','d'

和这样的字典（“my_dict”）： 'a':5, 'b':7', 'c':2, 'd':4

我想将这些结合起来创建一个 DataFrame，其中包含一个附加列，其中包含来自 my_dict 的相应值。

目前我正在使用以下方法，该方法适用于小型数据集，但效率非常低，并且会在我的完整数据集上导致 ***Error

import pyspark.sql.functions as F

# start with an arbitrary df containing "col1"
# initialise new column with zeros
my_df = my_df.withColumn('dict_data', F.lit(0))

for k,v in my_dict.items():
    my_df = my_df.withColumn('dict_data',
                             F.when((my_df['col1']==k),
                                     v).otherwise(df['dict_data'])
                             )

有没有更好的方法来做到这一点？我尝试过使用 Window 函数，但我很难在这种情况下应用它......

【参考方案1】：

您只需根据第一列的值将字典值映射到新列。可以参考：

pyspark create new column with mapping from a dict

【参考方案2】：

您可以使用中间数据框和连接来完成：

rows = ['col1': key, 'dict_data': value for key,value in my_dict.items()]
my_dict_df = rdd.parallelize(rows).toDF()

result_df = my_df.join(my_dict_df, 'col1', 'left')