如何使用列值作为 PySpark 中字典的键？答案

【问题标题】：How to use a column value as key to a dictionary in PySpark?如何使用列值作为 PySpark 中字典的键？
【发布时间】：2022-01-24 11:50:19
【问题描述】：

我有一个小的 PySpark DataFrame df:

index    col1
0        1    
1        3
2        4

还有一本字典：

LOOKUP = {0: 2, 1: 5, 2: 5, 3: 4, 4: 6}

我现在想在df 中添加一个额外的列col2，等于col1 的LOOKUP 值。

我的输出应该是这样的：

index    col1 col2
0        1    5    
1        3    4
2        4    6

我尝试使用：

df = df.withColumn(col("col2"), LOOKUP[col("col1")])

但这给了我错误，以及使用expr。

如何在 PySpark 中实现这一点？

【问题讨论】：

标签： python dataframe apache-spark pyspark apache-spark-sql

【解决方案1】：

您可以使用从 lookup 字典创建的 map 列：

from itertools import chain
from pyspark.sql import functions as F

lookup = {0: 2, 1: 5, 2: 5, 3: 4, 4: 6}
lookup_map = F.create_map(*[F.lit(x) for x in chain(*lookup.items())])

df1 = df.withColumn("col2", lookup_map[F.col("col1")])

df1.show()
#+-----+----+----+
#|index|col1|col2|
#+-----+----+----+
#|    0|   1|   5|
#|    1|   3|   4|
#|    2|   4|   6|
#+-----+----+----+

另一种方法是从字典中创建一个lookup_df，然后加入您的数据框

【讨论】：

【解决方案2】：

您可以在此处使用带有 python f-strings 的 CASE WHEN 语句和 LOOKUP 字典：

from pyspark.sql import functions as F
column = 'col1' #column to replace
e = f"""CASE {' '.join([f"WHEN {column}='{k}' THEN '{v}'" for k,v in LOOKUP.items()])} 
        ELSE NULL END"""
out = df.withColumn("col2",F.expr(e))

out.show()

+-----+----+----+
|index|col1|col2|
+-----+----+----+
|    0|   1|   5|
|    1|   3|   4|
|    2|   4|   6|
+-----+----+----+

【讨论】：