Pyspark应用pandasUDF后collect DataFrame到列表耗时过长如何优化？

耗时过高并非来自数据收集环节，而是Spark惰性求值机制导致pandas UDF被重复执行了两次：

• 第一次是df.show()触发全量计算
• 第二次是toPandas()/collect()触发了完整DAG的二次计算，相当于groupBy和pandas UDF逻辑跑了两遍，这才是耗时大幅上升的核心原因

1. 增加缓存避免重复计算 #

在pandas UDF计算完成后先缓存DataFrame，第一次action触发的计算结果会被存储，后续操作直接读取缓存即可，无需重新执行UDF。

2. 精简数据拉取逻辑 #

不需要拉取全量字段到Driver端，仅选择你需要的col_name字段即可，减少序列化、网络传输的额外开销。

df = df.groupBy("col_name").apply(pandas_udf_function)
# 缓存计算结果
df.cache()
# 第一次action触发计算，结果将被缓存
df.show()

# 仅拉取需要的字段生成列表
id_list = df.select("col_name").toPandas()["col_name"].tolist()
print(id_list)

# 用完后手动释放缓存
df.unpersist()

• 如果使用Spark 3.0+版本，可以调整spark.sql.execution.arrow.maxRecordsPerBatch参数，调大Arrow批次传输的大小，适配你的数据量提升pandas UDF执行效率
• 若数据量较大，可以优先用df.select("col_name").rdd.map(lambda x: x[0]).collect()的方式生成列表，内存开销比全量转pandas更低
• 排查pandas UDF内部逻辑是否存在可优化点，以及groupBy字段是否存在数据倾斜，进一步降低UDF本身的执行耗时

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