PySpark中使用Union、Cache、Checkpoint优化循环处理以解决磁盘空间不足问题

看起来你遇到了PySpark循环处理数据时的磁盘空间瓶颈，这种情况在循环Union操作中很常见，我来帮你拆解问题并给出优化方案：

问题根源分析 #

你当前的循环Union方式会让Spark维护一条不断增长的血统（Lineage）——每次循环都把新的Hmet_i和之前的Hmet合并，Spark需要追踪所有中间DF的依赖关系，导致临时shuffle文件越积越多，最终占满磁盘。另外，你使用的cache和checkpoint如果没有正确释放资源，也会加剧磁盘占用。

具体优化方案 #

1. 避免循环读取+Union，改用Spark原生批量读取

既然你是读取PATH下的所有子目录（bucket），完全可以让Spark一次性读取整个目录，省去循环的开销，也避免多次Union：

# 直接读取整个目录下的所有bucket，Spark会自动识别子目录结构
hda_all = spark.read.load(PATH).dropDuplicates().cache()

# 一次性完成join和计算
df_jpn_all = df_original.join(hda_all, on=['user', 'date'], how='inner')
Hmet = apply_function(df_jpn_all).dropDuplicates()

这种方式从根源上消除了循环Union带来的血统膨胀问题，效率会高很多。

2. 正确使用Cache和Checkpoint（如果必须循环的话）

如果因为特殊原因必须保留循环逻辑，一定要注意资源释放：

• Cache后及时Unpersist：每次循环处理完hda_i后，调用hda_i.unpersist()释放缓存，避免多个hda_i占用内存/磁盘；
• 设置独立的Checkpoint目录：默认的Checkpoint目录可能在临时磁盘，要指定一个空间充足的路径，并且Checkpoint后释放原DF：

# 先设置全局Checkpoint目录（在循环外执行一次即可）
spark.sparkContext.setCheckpointDir("/path/to/large-disk/checkpoint")

Hmet = None
for i, user_group in enumerate(os.listdir(PATH)):
    hda_i = spark.read.load(os.path.join(PATH, user_group)).dropDuplicates().cache()
    
    df_jpn = df_original.join(hda_i, on=['user', 'date'], how='inner')
    Hmet_i = apply_function(df_jpn).dropDuplicates().checkpoint()
    
    # 释放当前hda_i的缓存，避免累积
    hda_i.unpersist()
    
    if Hmet is None:
        Hmet = Hmet_i
    else:
        # 推荐用unionByName，避免列顺序不一致导致的问题
        Hmet = Hmet.unionByName(Hmet_i)
    
    # 可选：每次Union后也Checkpoint Hmet，截断血统
    Hmet = Hmet.checkpoint()

3. 优化Join操作，减少Shuffle数据

• 如果df_original数据量较小，使用广播变量减少Shuffle：

from pyspark.sql.functions import broadcast

df_jpn = broadcast(df_original).join(hda_i, on=['user', 'date'], how='inner')

• 检查user和date作为Join键是否存在数据倾斜，如果有，可以考虑加盐分区等方式优化。

4. 写入阶段的磁盘优化

• 调整输出分区数，避免生成大量小文件（小文件会浪费磁盘空间且影响读取性能）：

# 根据数据量设置合理的分区数，比如100
Hmet.repartition(100).write.format('parquet').mode('overwrite').save(os.path.join(PATH, fname))
# 或者用coalesce减少分区（适合数据量不大的情况）
Hmet.coalesce(20).write.format('parquet').mode('overwrite').save(os.path.join(PATH, fname))

• 检查Spark临时目录：修改spark.local.dir配置，将临时Shuffle文件指向空间更大的磁盘，在Spark初始化时设置：

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .config("spark.local.dir", "/path/to/large-disk/tmp") \
    .getOrCreate()

额外排查建议 #

• 查看Spark UI的Storage和Jobs页面，追踪哪些DF占用了过多磁盘/内存；
• 检查集群的磁盘使用情况，确认是临时Shuffle文件还是Checkpoint文件占满磁盘；
• 避免在循环中频繁创建DF，尽量复用或合并操作，减少中间文件生成。

备注：内容来源于stack exchange，提问作者s223