如何处理Spark中的不均匀分区？S3输出分区大小失衡问题

针对你遇到的300GB DataFrame按DataPartition, PartitionYear, PartitionStatement分区后数据严重不均（40GB vs 10MB），且直接用repartition(100)会生成大量小文件导致后续Spark作业异常的问题，我分享几个生产环境中验证过的解决方案：

方案1：拆分处理大分区+保留小分区完整性 #

核心思路是把大分区和小分区分开处理，避免小分区被强制拆分：

1. 先通过近似统计（避免全量扫描）筛选出数据量超过阈值（比如1GB）的大分区列表
2. 对大分区的数据单独调整分区数（用coalesce避免shuffle，或repartition按需拆分）
3. 小分区的数据保持原结构直接输出
4. 合并两部分数据后按原业务分区字段写入S3

示例代码：

// 用近似统计筛选大分区（避免全量count耗时）
val partitionStats = df.groupBy("DataPartition", "PartitionYear", "PartitionStatement")
  .countApprox(1000L, 0.95) // 近似统计，误差5%以内
  .filter($"count" > 1024*1024*1024) // 筛选大于1GB的分区
  .select("DataPartition", "PartitionYear", "PartitionStatement")
  .collect()
  .map(row => (row.getString(0), row.getString(1), row.getString(2)))
  .toSet

// 处理大分区：按500MB/文件估算，设置对应分区数
val largePartDf = df.filter(
  $"DataPartition".isin(partitionStats.map(_._1).toSeq: _*) &&
  $"PartitionYear".isin(partitionStats.map(_._2).toSeq: _*) &&
  $"PartitionStatement".isin(partitionStats.map(_._3).toSeq: _*)
).coalesce(400) // 假设大分区总数据约200GB，500MB/文件需400个分区

// 处理小分区：直接保留原分区逻辑
val smallPartDf = df.filter(
  !($"DataPartition".isin(partitionStats.map(_._1).toSeq: _*) &&
  $"PartitionYear".isin(partitionStats.map(_._2).toSeq: _*) &&
  $"PartitionStatement".isin(partitionStats.map(_._3).toSeq: _*))
)

// 合并后按业务分区输出
largePartDf.union(smallPartDf)
  .write
  .partitionBy("DataPartition", "PartitionYear", "PartitionStatement")
  .parquet("s3://your-bucket/target-path")

方案2：利用Spark参数自动控制文件大小 #

Spark内置参数可以帮你自动平衡文件大小，无需手动拆分分区：

• spark.sql.files.maxPartitionBytes：单个文件最大字节数（默认1GB）
• spark.sql.files.maxRecordsPerFile：单个文件最大记录数（优先级高于上一个参数）

你可以在输出前临时设置参数，让大分区自动拆分、小分区保持单文件：

// 临时设置单个文件最大为500MB
spark.conf.set("spark.sql.files.maxPartitionBytes", "512m")

df.write
  .partitionBy("DataPartition", "PartitionYear", "PartitionStatement")
  .parquet("s3://your-bucket/target-path")

这个方案最简单，适合不想修改分区逻辑的场景，注意不要把阈值设得太小（会生成过多文件）或太大（单个文件过大导致读取内存压力）。

方案3：调整分区粒度，拆分大分区字段 #

如果业务允许，可以对大分区对应的字段做更细粒度拆分：

• 比如PartitionYear是年份，大分区可以新增PartitionMonth子分区
• 或者对PartitionStatement做哈希拆分（比如取前两位字符作为子分区）

这样原本40GB的大分区会被拆分成多个均衡的子分区，小分区保持原粒度，不会生成过多小文件：

df.withColumn("PartitionMonth", substring($"PartitionDate", 6, 2)) // 假设存在日期字段PartitionDate
  .write
  .partitionBy("DataPartition", "PartitionYear", "PartitionMonth", "PartitionStatement")
  .parquet("s3://your-bucket/target-path")

方案4：结合业务字段的全局重分区 #

如果集群资源充足，可以用repartition结合业务字段调整分区数，既保证数据均衡又避免小文件：

df.repartition(
  2000, // 总分区数接近原分区数，保证每个业务分区内数据均衡
  $"DataPartition", $"PartitionYear", $"PartitionStatement"
)
.write
.partitionBy("DataPartition", "PartitionYear", "PartitionStatement")
.parquet("s3://your-bucket/target-path")

这个方法会触发shuffle，但Spark会先按业务字段分组，再在组内拆分分区，避免小分区被强制拆分成多个文件。

避坑提醒 #

• 不要直接用repartition(100)全局重分区：会完全打乱业务分区局部性，触发大量shuffle，且小分区会被拆成过多小文件
• 小文件过多会导致Spark读取时生成数千个任务，占用过多资源甚至触发集群异常
• 统计分区数据量优先用近似方法：全量count会扫描300GB数据，耗时极长

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