如何在PySpark中写入DataFrame时手动控制单个Part-file大小？

以下是几种实用的手动控制单文件大小的方案，可根据数据特点选择：

1. 基于数据量估算调整分区数 #

这是最直接的方法：先估算DataFrame的总数据量，结合目标单文件大小计算所需分区数，再通过repartition或coalesce调整分区后写入。

步骤： #

• 估算DataFrame的近似总大小（字节）
• 根据目标单文件大小计算分区数（分区数 = 总大小 / 目标大小，确保至少为1）
• 调整分区并写入

代码示例： #

# 估算DataFrame的近似大小（遍历每个分区计算单条记录字符串长度之和，仅作参考）
total_size_bytes = df.rdd.mapPartitions(lambda part: [sum(len(str(row)) for row in part)]).sum()

# 设置目标单文件大小（示例为1GB）
target_file_size = 1 * 1024 * 1024 * 1024  # 1GB

# 计算所需分区数
num_partitions = max(1, int(total_size_bytes / target_file_size))

# 调整分区：如果是减少分区用coalesce（无shuffle，效率更高）；需要重新分配用repartition（全shuffle）
df = df.coalesce(num_partitions) if num_partitions < df.rdd.getNumPartitions() else df.repartition(num_partitions)

# 写入文件
df.write.parquet("/path/to/output/dir")

注意：如果使用压缩格式（如Snappy、Gzip），需结合压缩比调整目标大小（比如Snappy压缩比约3-5倍，目标大小需乘以压缩比来估算原始数据量）。

2. 限制每个文件的最大记录数 #

通过Spark配置spark.sql.files.maxRecordsPerFile，强制每个输出文件的记录数不超过设定值，间接控制文件大小（适合记录大小相对均匀的场景）。

代码示例： #

# 设置每个文件最多100万条记录
spark.conf.set("spark.sql.files.maxRecordsPerFile", 1000000)

# 写入时自动按记录数拆分文件
df.write.parquet("/path/to/output/dir")

局限性：如果记录大小差异极大（比如有的记录几KB，有的几MB），文件大小会出现明显波动。

3. 自定义分区/分桶逻辑 #

针对数据分布有规律的场景，可通过自定义分区规则（如数值范围、时间区间）或分桶，让每个分区的数据量接近目标大小。

示例：按数值范围手动分区 #

假设数据有一个数值型字段user_id，可按user_id的范围拆分，确保每个区间的数据量均匀：

# 获取user_id的极值
min_uid = df.selectExpr("min(user_id)").first()[0]
max_uid = df.selectExpr("max(user_id)").first()[0]
uid_range = max_uid - min_uid

# 计算每个分区的user_id步长
num_partitions = max(1, int(total_size_bytes / target_file_size))
step = uid_range // num_partitions

# 生成分区过滤条件
partition_conditions = []
for i in range(num_partitions):
    start = min_uid + i * step
    end = min_uid + (i + 1) * step
    if i == num_partitions - 1:
        # 最后一个分区包含最大值
        condition = f"user_id >= {start}"
    else:
        condition = f"user_id >= {start} AND user_id < {end}"
    partition_conditions.append(condition)

# 逐个写入分区
for idx, cond in enumerate(partition_conditions):
    partition_df = df.filter(cond)
    partition_df.write.mode("append").parquet(f"/path/to/output/partition={idx}")

分桶方案（适合需要后续查询优化的场景） #

通过bucketBy指定分桶数和分桶字段，让数据均匀分布到各个桶中，每个桶对应一个输出文件：

# 按user_id分桶，设置分桶数为计算出的num_partitions
df.write.bucketBy(num_partitions, "user_id").sortBy("user_id").parquet("/path/to/output/dir")

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