PySpark基于topic_idx关联DataFrame添加列，需保留原行数无重复

核心思路 #

利用df2中topic_idx与cities的唯一映射关系，先将df2转换为字典，再通过映射的方式为df1添加cities列，完全保证df1的行数和原有内容不变。

Pandas 实现方案 #

1. 将df2转换为topic_idx到cities的字典

# 基于df2构建映射字典
topic_cities_map = df2.set_index('topic_idx')['cities'].to_dict()

2. 为df1添加cities列

# 直接通过map方法完成映射，无重复行
df1['cities'] = df1['topic_idx'].map(topic_cities_map)

• 若df1中存在df2没有的topic_idx，对应cities值会显示NaN，可根据需求替换为默认值（比如空列表）。

Spark 实现方案（含自定义UDF） #

Spark分布式环境下需用广播变量优化性能，结合UDF实现：

1. 构建映射字典并广播

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

# 从df2收集数据生成字典
topic_cities_map = {row.topic_idx: row.cities for row in df2.collect()}
# 广播字典，避免每个Task重复加载
broadcast_map = spark.sparkContext.broadcast(topic_cities_map)

2. 定义并使用UDF添加列

# 自定义映射函数
def get_cities(topic_idx):
    return broadcast_map.value.get(topic_idx, [])  # 无匹配时返回空列表，可按需修改

# 注册UDF（指定返回类型为字符串数组，对应cities的列表类型）
get_cities_udf = udf(get_cities, ArrayType(StringType()))

# 为df1添加cities列
df1 = df1.withColumn('cities', get_cities_udf(df1['topic_idx']))

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