PySpark UDF中迭代DataFrame列对象写入字典失败求助

问题根源 #

• 惰性执行未触发计算：PySpark的withColumn仅生成执行计划，不会立刻执行UDF代码。只有调用show()、count()、collect()这类动作（Action）操作时，才会实际触发UDF运行，否则字典永远不会被修改。
• Driver与Executor内存隔离：即便触发计算，Driver端定义的dictionary和Executor中UDF访问的字典是完全独立的内存副本。UDF修改的是Executor本地的字典，Driver端的原字典无法同步这些修改，最终仍为空。

解决方案 #

小数据集直接处理 #

如果数据量不大，直接将数据拉取到Driver端遍历构建字典，简单高效：

dictionary = {}
# 把需要的列数据收集到Driver
rows = df.select("valueA", "valueB", "valueC").collect()
# 遍历每一行填充字典
for row in rows:
    x, y, z = row.valueA, row.valueB, row.valueC
    if x not in dictionary:
        dictionary[x] = {z: {"data": y}}
    else:
        dictionary[x].update({z: {"data": y}})

大数据集分布式处理 #

数据量较大时，先通过Spark分布式聚合操作处理数据，再将结果拉取到Driver构建字典：

from pyspark.sql import functions as F

# 按valueA分组，把每组的(valueC, valueB)聚合成映射结构
aggregated_df = df.groupBy("valueA").agg(
    F.map_from_entries(
        F.collect_list(F.struct("valueC", F.struct(F.col("valueB").alias("data"))))
    ).alias("key_mapping")
)

# 收集聚合结果到Driver，生成目标字典
dictionary = {row.valueA: row.key_mapping for row in aggregated_df.collect()}

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