Spark如何实现单客户单输出文件且避免OOM？

针对你这个场景，刚好有几个合适的解决方案，既能保证每个(client_id, report_date)分区只生成单个文件，又不会触发OOM——毕竟你提到单个客户文件体积不大：

1. 精准按目标分区字段重分区（最简洁方案） #

你之前的问题出在repartition(2, "customer_id", "report_date")手动指定了分区数2，导致每个(client_id, report_date)组合被拆分到2个分区里。其实只要直接把client_id和report_date作为repartition的分区依据，Spark会自动为每一组唯一的(client_id, report_date)分配一个单独的分区，后续写入时每个分区就只会生成一个文件。

修正后的代码（注意你原代码里的customer_id应该是client_id的笔误，value1对应value_1）：

df.repartition(col("client_id"), col("report_date"))
  .sortWithinPartitions("date", "value_1")
  .write.partitionBy("client_id", "report_date")
  .csv("results")

原理：repartition按这两个字段分区后，每个分区内只有同一客户、同一报告日期的数据，后续partitionBy只是将数据写入对应的目录，每个目录自然只会有一个文件。因为单个客户文件体积小，即使有大量客户，也不会出现OOM——每个分区的数据量都在可控范围内。

2. 自定义分组写入（更灵活的控制） #

如果你需要更精细的写入控制（比如自定义文件名、处理特殊情况），可以通过groupBy结合foreachPartition实现，直接对每个(client_id, report_date)分组的数据单独写入：

import org.apache.spark.sql.{Row, SaveMode}
import org.apache.spark.sql.functions._

df.groupBy("client_id", "report_date")
  .foreachPartition { partitionIter =>
    if (partitionIter.nonEmpty) {
      // 取出分组标识（注意数据类型要和你的DataFrame匹配）
      val firstRow = partitionIter.next()
      val clientId = firstRow.getAs[String]("client_id")
      val reportDate = firstRow.getAs[String]("report_date")
      
      // 将当前分区的数据转换为DataFrame并排序
      val groupData = spark.createDataFrame(partitionIter ++ Iterator(firstRow), df.schema)
        .sort("date", "value_1")
      
      // 构建目标路径
      val outputDir = s"results/client_id=$clientId/report_date=$reportDate"
      
      // 写入文件，可根据需求设置保存模式
      groupData.write.mode(SaveMode.Overwrite).csv(outputDir)
    }
  }

这个方法的优势是完全掌控每个分组的写入逻辑，确保每个分区只生成一个文件，同时因为每个分组数据量小，不会有内存溢出的风险。

注意事项 #

• 确保client_id和report_date的数据类型一致（比如都是字符串或日期类型），避免因类型不匹配导致分区混乱。
• 如果你的集群资源有限，且存在超大量的(client_id, report_date)组合，第一种方案的repartition会创建对应数量的分区，此时可以适当限制并发度（比如通过spark.sql.shuffle.partitions调整），但因为每个分区数据量小，通常不会有问题。

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