Spark FlatMap处理大XML生成海量结果写入Parquet失败的解决方法

我用flatMap函数将超大XML文件拆分为数万个小型XML字符串片段，并尝试写入Parquet格式，但操作的Stage失败率极高，推测是DataFrameWriter写入阶段丢失Executor，可能因超出存储限制。

我原本以为实现Iterable类型的flatMap可让Spark逐个拉取结果，避免内存缓冲，但堆栈跟踪显示Spark会在写入前将每个flatMap的所有结果存入数组（堆栈信息：at scala.collection.mutable.ResizableArray.foreach(ResizableArray.scala:62)），这与预期不符。

以下是flatMap中使用的类的伪代码，该类返回Iterable类型：

class XmlIterator(filepath: String, split_element: String) extends Iterable[String] {

   // 基于文件路径的FileInputStream打开XMLEventReader
   // 实现一个Iterable，每次返回XML文件的一个片段

   def iterator = new Iterator[String] {
      def hasNext = { 
        // 推进输入流，若有可返回内容则返回true
      }
      def next = {
        // 返回当前XML片段字符串
      }
  }
}

使用方式如下：

var dat = [包含大量GB级文件路径的单列DataFrame]

dat.repartition(1375) // 按行数重分区，希望DataFrameWriter处理完每个文件就立即写入
  .flatMap(rec => new XmlIterator(rec, "bibrecord"))
  .write
  .parquet("some_path")

少量文件并行处理正常，但大批次处理会出现Stage失败。请问有没有更内存高效的替代策略来保存flatMap的结果？

1. 改用mapPartitions替代flatMap，直接操作迭代器 #

Spark的flatMap处理Iterable时，底层会先将整个Iterable的内容收集到数组中再处理，这就是你看到ResizableArray堆栈信息的原因。换成mapPartitions可以直接返回迭代器，让Spark逐个拉取元素，避免一次性加载所有XML片段到Executor内存：

dat.repartition(1375)
  .mapPartitions { iter =>
    iter.flatMap { rec =>
      new XmlIterator(rec, "bibrecord").iterator
    }
  }
  .write
  .parquet("some_path")

2. 优化Spark内存配置，缓解OOM压力 #

• 调高spark.executor.memory，给每个Executor分配足够内存处理大文件的XML解析
• 调低spark.executor.cores，减少单个Executor同时处理的任务数，降低内存竞争
• 开启spark.memory.offHeap.enabled并配置spark.memory.offHeap.size，利用堆外内存存储临时数据，减轻堆内存压力

3. 拆分任务批次，降低单任务负载 #

如果单个文件拆分出的XML片段过多，即使改用迭代器也可能因任务输出数据量过大导致内存溢出。可以将文件路径DataFrame拆分为多个小批次处理：

// 示例：分批次处理文件
val filePaths = dat.collect().map(_.getString(0))
filePaths.grouped(100).foreach { batch =>
  spark.createDataset(batch)
    .mapPartitions { iter =>
      iter.flatMap(path => new XmlIterator(path, "bibrecord").iterator)
    }
    .write.mode("append").parquet("some_path")
}

这种方式将大任务拆分为多个小任务，每个批次处理的文件更少，降低单个Executor的内存负载。

4. 使用Spark原生XML解析库替代自定义实现 #

Spark官方的spark-xml库（com.databricks:spark-xml_2.12:0.17.0）支持直接按指定元素拆分大XML文件，底层已做流式解析优化，无需自行实现迭代器：

import com.databricks.spark.xml._

spark.read
  .option("rowTag", "bibrecord") // 指定要拆分的目标元素
  .xml("path_to_large_xml_files")
  .write.parquet("some_path")

该库会流式读取XML内容，不会一次性加载整个文件到内存，稳定性和效率优于自定义实现。

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