Spark DataFrame转换数组后导出CSV或转Pandas超时如何解决？

问题诱因 #

• Spark的惰性计算特性：你观测到DataFrame状态正常仅代表Spark生成了合法的逻辑执行计划，并未实际触发计算。调用toPandas()、write.csv()这类Action操作时，才会真正执行全量50万条数据的UDF转换逻辑，此时才会暴露执行超时问题。
• 自定义Python UDF的额外开销：Python UDF需要在JVM（Spark执行层）和Python进程之间反复序列化、传输数据，你需要转换的embeddings属于高维大数组，单条数据量本身偏大，50万条总数据量远超默认的socket通信超时阈值。
• 多字段重复UDF调用放大开销：你同时对5个字段应用了自定义UDF，相当于每条数据要完成5次跨进程数据传输和计算，进一步拖慢执行速度，最终触发超时。

解决方法 #

方案1：替换自定义Python UDF为Spark原生函数（优先推荐） #

Spark原生内置的array_join、concat函数可以直接实现数组转指定格式字符串的需求，完全运行在JVM层，无跨进程通信开销，性能比Python UDF高10倍以上，代码示例如下：

from pyspark.sql.functions import concat, lit, array_join, col

result = result \
    .withColumn("token", concat(lit("["), array_join(col("token"), ","), lit("]"))) \
    .withColumn("ner", concat(lit("["), array_join(col("ner"), ","), lit("]"))) \
    .withColumn("embeddings", concat(lit("["), array_join(col("embeddings"), ","), lit("]"))) \
    .withColumn("ner_chunk", concat(lit("["), array_join(col("ner_chunk"), ","), lit("]"))) \
    .withColumn("document", concat(lit("["), array_join(col("document"), ","), lit("]")))

方案2：调整Spark网络超时配置（仅用于必须使用Python UDF的场景） #

如果业务逻辑必须依赖Python UDF实现，可以在初始化SparkSession时调大网络超时阈值，同时开启PySpark Arrow优化提升序列化效率，配置示例如下：

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("TweetProcess") \
    .config("spark.executor.heartbeatInterval", "120s") \
    .config("spark.network.timeout", "600s") \
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    .getOrCreate()

方案3：优化全量数据拉取逻辑 #

如果需要导出全量数据，避免直接调用toPandas()把所有数据拉到Driver端，建议直接用Spark原生的write.csv()方法分布式写入文件；如果确实需要转Pandas，可以先对数据做分片，分批拉取转换。

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