Spark技术问题：如何收集大量数据而不出现内存溢出

阿华AIGC实验室

2026-5-27

解决Spark查询Parquet后Collect内存溢出的问题

我之前也踩过这个一模一样的坑！当你用collect()把shuffle后的全量结果拉到Driver内存时，数据量一大必然会OOM——毕竟Driver的内存资源本来就有限，根本扛不住全量的分布式结果。你提到的按城市列分片处理的思路非常靠谱，我来补充几个优化细节和更稳妥的实现方式：

一、优先用分布式分区替代Driver端循环

其实不用手动遍历城市列表，Spark本身就支持按指定字段分区后直接写入磁盘，全程分布式处理，完全不用把数据拉到Driver：

// 先执行你的原始查询得到DataFrame
val df = spark.sql("你的原始SQL查询语句")

// 按city字段分区，自动把同城市的数据分到一起，然后写入磁盘
df.repartition($"city")
  .write
  .mode("overwrite")
  .parquet("/hdfs/path/to/save/result_by_city")

这样Spark会自动帮你把数据按城市分片存储，每个分片的写入操作由Executor完成，Driver只负责调度，完全不会出现OOM问题，效率也比手动循环高得多。

二、如果必须手动遍历城市的优化方案

要是你确实需要在Driver端对每个城市的结果做特殊处理，那可以按以下步骤来，避免踩坑：

1. 安全获取去重城市列表

先拿到所有去重的城市值，但注意不要直接collect大列表——如果城市数量也很多，这一步就先OOM了：

// 如果城市数量不多，可以直接collect
val cities = df.select("city").distinct().collect().map(_.getString(0))

// 如果城市数量极大，先把去重结果写到磁盘再读取，避免占Driver内存
df.select("city").distinct()
  .repartition(1)
  .write
  .mode("overwrite")
  .text("/tmp/distinct_cities")
// 然后从磁盘读取城市列表（示例为Scala读取本地文件，实际可根据存储路径调整）
val cities = scala.io.Source.fromFile("/tmp/distinct_cities/part-00000").getLines().toArray

2. 参数化查询避免SQL注入&语法错误

直接字符串拼接WHERE city = '$city'很容易因为城市名里的特殊字符（比如单引号）报错，甚至有SQL注入风险，改用DataFrame API或者参数化SQL更安全：

for (city <- cities) {
  // 方式1：用DataFrame API过滤，更简洁安全
  val cityDF = df.filter($"city" === city)
  
  // 方式2：参数化SQL查询
  // val cityDF = spark.sql("SELECT * FROM t WHERE city = ?", city)
  
  // 保存到磁盘，或者做你需要的本地处理
  cityDF.write.mode("overwrite").parquet(s"/hdfs/path/to/save/result_$city")
  
  // 非必要不要用collect！如果必须要拉到Driver，确保单城市数据量不大
  // val localData = cityDF.collect()
  // ... 你的本地处理逻辑
}

3. 把处理逻辑放到Executor端

如果你的处理逻辑不需要在Driver完成，尽量用foreachPartition把逻辑放到Executor执行，完全避免数据拉到Driver：

df.foreachPartition { partitionIter =>
  // 这段代码在每个Executor的分区上执行，不需要Driver参与
  // 示例：自定义写文件逻辑
  val writer = new java.io.PrintWriter(new java.io.File(s"/local/path/result_${Thread.currentThread().getId}.txt"))
  partitionIter.foreach(row => writer.println(row.mkString(",")))
  writer.close()
}