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

标签： pythonpyspark

Answer 1

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

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