Spark如何在Driver实例内存中迭代RDD写入不支持Sink数据库

Spark 提供了完全匹配该需求的API，无需修改数据库客户端、也不需要全量拉取数据即可在Driver端迭代处理全量RDD/Dataset数据。

核心可用API #

直接使用RDD.toLocalIterator（Dataset/DataFrame可直接调用同名方法）即可，该方法的执行逻辑和collect()有本质区别：

• 不会一次性将全量数据拉取到Driver内存，而是按分区顺序依次拉取：每次仅将1个分区的数据传输到Driver端，迭代处理完当前分区的所有记录后，再拉取下一个分区的数据
• 整个迭代过程Driver端的内存峰值仅需要容纳单个分区的大小，只要单分区数据量不超过Driver可用内存，就不会出现OOM问题
• 迭代逻辑全程运行在Driver进程内，不需要将数据库连接对象序列化分发到Executor节点，完全适配不可序列化的数据库通信对象场景。

适配场景的使用要点 #

• 提前调整分区大小：如果读取Parquet后单分区数据量过大，可以通过repartition方法重新调整分区数，将单分区数据量控制在Driver可承载的范围内
• 采用批量写入逻辑：不要逐条提交数据库写入请求，在迭代过程中攒够固定批量（根据单条数据大小调整，通常1000~10000条即可）后执行一次批量写入，大幅提升写入效率
• 注意资源保活：整个迭代过程中Spark会持续持有对应Executor的计算资源，直到所有分区遍历完成，需要合理设置作业超时时间，避免长写入过程被集群强制回收资源

示例代码 #

import org.apache.spark.sql.{Row, SparkSession}
import java.util.ArrayList

object WriteToCustomSink {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().getOrCreate()
    // 读取Parquet文件
    val sourceDF = spark.read.parquet("/path/to/your/parquet/directory")
    // Driver端初始化不可序列化的数据库客户端
    val dbClient = new CustomNonSerializableDbClient()
    val batchSize = 3000
    val writeBuffer = new ArrayList[Row](batchSize)

    try {
      val dataIterator = sourceDF.toLocalIterator
      while (dataIterator.hasNext) {
        val currentRow = dataIterator.next()
        writeBuffer.add(currentRow)
        // 达到批量阈值执行写入
        if (writeBuffer.size() >= batchSize) {
          dbClient.batchInsert(writeBuffer)
          writeBuffer.clear()
        }
      }
      // 写入最后剩余的不足批量的数据
      if (!writeBuffer.isEmpty) {
        dbClient.batchInsert(writeBuffer)
      }
    } finally {
      // 确保资源释放
      dbClient.close()
      spark.stop()
    }
  }
}

方案对比 #

• 不推荐使用collect()：会全量拉取所有数据到Driver内存，数据量超过内存阈值时直接触发OOM
• 不推荐使用foreach/foreachPartition：算子逻辑运行在Executor端，要求传入的数据库客户端必须支持序列化，和当前场景约束冲突

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