PySpark DataFrame数组元素映射至另一DataFrame的实现咨询

有两种常用实现方案，都可以保证数组原有顺序不变，满足和column3按索引对应的要求：

方案1：小映射表场景（df_two数据量小，推荐优先使用） #

核心思路是把df_two转成广播字典，直接遍历数组元素做替换，执行性能更高。

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

# 1、收集df_two的映射关系为字典
id_name_map = df_two.rdd.map(lambda x: (x["columnA"], x["columnB"])).collectAsMap()
# 2、把映射字典广播到所有执行节点，减少内存冗余
broadcast_map = spark.sparkContext.broadcast(id_name_map)

# 3、定义数组替换UDF
def map_array_elements(arr):
    return [broadcast_map.value.get(elem) for elem in arr]
map_udf = udf(map_array_elements, ArrayType(StringType()))

# 4、生成结果表
result_df = df_one.withColumn("newColumn", map_udf(col("column2"))) \
                  .select("column1", "newColumn", "column3")

适用场景：df_two行数在10w以下，全量加载到内存不会占用过多资源的场景

方案2：大映射表场景（df_two数据量大，无法全量加载到内存） #

核心思路是把column2数组带原始索引炸开，关联df_two之后按索引排序再重新聚合为数组，完全保证顺序和原数组一致。

from pyspark.sql.functions import posexplode, collect_list, col

# 1、带索引炸开column2数组，pos字段为元素在原数组的对应索引
df_exploded = df_one.select("column1", "column3", posexplode("column2").alias("pos", "columnA"))

# 2、和df_two关联获取字符串值
df_joined = df_exploded.join(df_two, on="columnA", how="left")

# 3、按原始索引排序后重新聚合为数组，保证顺序不变
result_df = df_joined.orderBy("pos") \
                     .groupBy("column1", "column3") \
                     .agg(collect_list("columnB").alias("newColumn")) \
                     .select("column1", "newColumn", "column3")

如果column2存在df_two中没有匹配的元素，left join会返回null，要自定义默认值可以在join后用coalesce函数处理columnB字段即可

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