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

Question

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

我有一个Spark DataFrame（“ my_df”），其中有一个列（“ col1”），其值是'a'，'b'，'c'，'d'

和像这样的字典（“ my_dict”）：{'a'：5，'b'：7'，'c'：2，'d'：4}

我想将它们结合起来以创建一个DataFrame，并在其中包含my_dict中对应值的附加列。

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

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函数，但是在这种情况下很难应用...

Answer 1

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

pyspark使用字典中的映射创建新列

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