如何在PySpark DataFrame中通过字典映射列值？

你遇到的问题很典型——在UDF里直接调用RDD的lookup方法之所以会触发SPARK-5063错误，是因为UDF运行在Executor节点上，而RDD的lookup是Driver端的操作，你不能在分布式执行的代码里直接引用Driver端的RDD对象，这会导致上下文不匹配。

下面给你两种可行的解决方案，分别适合不同场景：

方法一：使用广播变量（Broadcast Variable） #

广播变量可以把小字典分发到所有Executor节点，避免重复传输，适合映射关系数据量不大的情况：

1. 先把RDD形式的映射关系转换成Python字典，然后广播：

# 把RDD转成本地字典
map_dict = dict(dicts.collect())
# 广播字典
broadcast_map = sc.broadcast(map_dict)

2. 编写UDF使用广播变量里的映射：

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

def map_val(x):
    return broadcast_map.value.get(x, x)  # 找不到匹配值时返回原内容，可按需调整

myfun = udf(map_val, StringType())

# 生成映射列
df = df.withColumn('col1_map', myfun('col1'))
df = df.withColumn('col2_map', myfun('col2'))

df.show()

方法二：使用DataFrame Join操作（更Spark原生的方式） #

如果映射关系的数据量比较大，广播变量可能占用过多Executor内存，这时候可以把RDD转换成DataFrame，通过join实现映射，这种方式更适合大数据场景：

1. 把RDD转换成映射DataFrame：

map_df = dicts.toDF(['original_val', 'mapped_val'])
map_df.show()
# +-------------+-----------+
# |original_val|mapped_val|
# +-------------+-----------+
# |        india|        ind|
# |          usa|         us|
# |        japan|        jpn|
# |     uruguay|        urg|
# +-------------+-----------+

2. 分别对col1和col2进行关联，然后重命名列：

# 处理col1的映射
df = df.join(map_df, df.col1 == map_df.original_val, how='left') \
       .withColumnRenamed('mapped_val', 'col1_map') \
       .drop('original_val')

# 处理col2的映射
df = df.join(map_df, df.col2 == map_df.original_val, how='left') \
       .withColumnRenamed('mapped_val', 'col2_map') \
       .drop('original_val')

df.show()

运行后就能得到你期望的结果：

+-----+-------+--------+--------+
| col1| col2|col1_map|col2_map|
+-----+-------+--------+--------+
|india| japan| ind| jpn|
| usa|uruguay| us| urg|
+-----+-------+--------+--------+

内容的提问来源于stack exchange，提问作者YOLO