PySpark大数据集（25TB）下该重分区代码能否高效执行？

我是Spark新手，现在有一个存在分类子文件夹（比如按车型名称分类）的小数据集，每行都包含partitionID列。我的目标是用PySpark代码把数据按partitionID重分区，让对应分区的数据写入对应的文件夹。现在的代码在小数据集上运行正常，但担心在25TB级的超大数据集上没法高效执行，尤其怕迭代过程中Driver端会在单节点物化DataFrame，没法做到分布式处理。附上代码：

from pyspark.sql import SparkSession
import os
import sys
from pyspark.sql import functions as F

...

builder = SparkSession.builder
spark = builder.getOrCreate()

root_folder_path   = "<ROOT_PATH>"
output_folder_path = "<OUTPUT_PATH>"

# Get a list of categories inside the directory
folder_names = [folder for folder in os.listdir(root_folder_path) if os.path.isdir(os.path.join(root_folder_path, folder))]

# Iterate over the list of folder names (or categories)
for category_name in folder_names:
    # Read the entity parquet files
    entity_df = spark.read.parquet(root_folder_path + category_name, header=True, inferSchema=True)
    
    # Iterate over distinct partition ID values and write rows to separate folders
    distinct_data_partition_ids = entity_df.select("partitionId").distinct().rdd.map(lambda x: x[0]).collect()

    for data_partition_id in distinct_data_partition_ids:
        # Filter rows for the current dataPartitionId
        filtered_df = entity_df.filter(entity_df["partitionId"] == data_partition_id)

        # Construct the output directory for this partition ID
        output_directory = os.path.join(output_folder_path + category_name + "/" + str(data_partition_id))
    
        # Write rows to the corresponding output directory
        filtered_df.write.parquet(output_directory, mode="overwrite")

现有代码的核心问题 #

• Driver端严重瓶颈：collect()会把所有partitionId的结果拉到Driver节点内存中，25TB数据的partitionId基数如果很大，直接会撑爆Driver内存；而且后续嵌套循环完全在Driver单节点串行执行，完全浪费Spark的分布式能力。
• 重复扫描数据：每个partitionId循环都会重新扫描整个entity_df做过滤，相当于每个分区都要读一遍原始数据，25TB量级下会把IO资源耗尽，性能暴跌。
• 小文件爆炸风险：每个partitionId单独写入，极易生成大量极小的Parquet文件，后续读取时会严重拖慢性能。

高效优化方案 #

直接用Spark内置的分区写入机制，完全避免Driver端串行迭代，充分利用分布式处理能力：

优化后代码 #

from pyspark.sql import SparkSession
import os
from pyspark.sql import functions as F

builder = SparkSession.builder.getOrCreate()
spark = builder.getOrCreate()

root_folder_path = "<ROOT_PATH>"
output_folder_path = "<OUTPUT_PATH>"

# 直接读取所有分类文件夹的数据，Spark支持路径通配符
entity_df = spark.read.parquet(os.path.join(root_folder_path, "*"))

# 从文件路径中解析分类名称，替代手动遍历文件夹的逻辑（索引根据实际路径调整）
entity_df = entity_df.withColumn(
    "category_name",
    F.element_at(F.split(F.input_file_name(), "/"), -2)
)

# 按分类和partitionId分布式分区写入，Spark自动处理所有并行逻辑
entity_df.write.mode("overwrite") \
    .partitionBy("category_name", "partitionId") \
    .parquet(output_folder_path)

优化点说明 #

• 全流程分布式处理：partitionBy会让Spark自动将数据在Executor节点间分区处理，完全不需要Driver端串行迭代，彻底规避单节点瓶颈。
• 仅扫描一次数据：整个流程只读取一次原始数据，避免重复IO操作，极大节省存储和网络资源。
• 控制小文件数量：可通过调整spark.sql.shuffle.partitions（默认200，根据集群规模调整）、spark.sql.files.maxRecordsPerFile参数，控制输出文件的大小和数量，避免小文件问题。
• 自动维护目录结构：partitionBy会自动生成category_name=xxx/partitionId=xxx的标准分区目录，和你手动创建的目录逻辑一致，且后续读取时Spark能自动识别分区信息，提升查询性能。

额外注意事项 #

• 如果必须去掉目录名前缀（比如不需要category_name=），可以在写入后用文件系统工具批量重命名，但优先推荐使用Spark标准分区格式，后续维护和读取性能更好。
• 针对25TB级数据集，建议开启动态分区（默认已开启），并设置spark.sql.sources.partitionOverwriteMode=dynamic，避免全量覆盖时的性能损耗。
• 提前确认partitionId的基数，如果基数过大（比如百万级以上），需要评估集群的存储和元数据处理能力，避免分区过多导致的压力。

内容的提问来源于stack exchange，提问作者snehil suresh wakchaure