PySpark中持久化DataFrame后两次写入仍重计算，如何优化？

核心问题 #

persist() 是惰性操作，仅标记DataFrame需要缓存，但不会立即执行计算。两次write都是action操作，第一次会触发计算并尝试缓存，但如果缓存级别不合适（比如默认的MEMORY_ONLY，数据超出内存时会被自动丢弃），第二次write仍会重新计算。

调整方法 #

1. 强制触发缓存生效：在persist()后执行一个action操作（比如count()），确保DataFrame的计算结果被实际缓存到存储介质中。
2. 指定可靠的缓存级别：根据数据量选择MEMORY_AND_DISK这类级别，内存存不下的部分会写入磁盘，避免缓存失效。

修改后的代码示例：

from pyspark.storagelevel import StorageLevel

# 指定MEMORY_AND_DISK级别缓存，兼顾内存与磁盘存储
df.persist(StorageLevel.MEMORY_AND_DISK)
# 执行count()触发计算，将数据实际缓存下来
df.count()

# 两次写入操作直接读取缓存，无需重复计算
df.write.mode("overwrite").csv(file_path_1)
df.write.mode("overwrite").csv(file_path_2)

# 不再需要缓存时释放资源
df.unpersist()

替代方案：复用缓存后的RDD #

如果场景允许，也可以先缓存DataFrame对应的RDD，再基于RDD生成DataFrame完成写入：

from pyspark.storagelevel import StorageLevel

rdd = df.rdd.persist(StorageLevel.MEMORY_AND_DISK)
rdd.count()

# 基于缓存的RDD生成DataFrame后写入
rdd.toDF(df.schema).write.mode("overwrite").csv(file_path_1)
rdd.toDF(df.schema).write.mode("overwrite").csv(file_path_2)

rdd.unpersist()

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