[Spark] OOM问题及解决方案

1. 背景

最近在业务开发过程中，遇到如下需求：

一张Hive表中存储着item id和描述这个id的文本（已经切词，各个词语之间' '分隔）。另外还有一份数据，其中存储了各个词语和该词语对应的embedding vector。现要计算每个id对应文本的词向量表示，即将同一个id对应的文本中所有词语embedding vector求和。

2. 问题描述

在计算embedding vector求和过程中，出现了OOM问题。

//代码1
val sql = s"select id, text from table where day = '2020-08-20'"

val wordToVectorMap: collection.Map[String, Array[Float]] = ... // 数据量大约2G
val EMBEDDING_SIZE = 300
val zeroVector = new Array[Float](EMBEDDING_SIZE)
val idToWordVector: DataFrame = hiveContext.sql(sql).rdd.map(x => {
  
  val id = x.getString(0)
  val wordList: Array[String] = x.getString(1).split(" ")

  var result: Array[Float] = zeroVector
  for (word <- wordList) {
    if (wordVectorMap.contains(word)) {
      val vec = wordVectorMap.getOrElse(word, zeroVector) // OOM
      for (i <- 0 until EMBEDDING_SIZE) {
        result(i) = result(i) + vec(i)
      }
    }
  }
  (id, result.mkString(",")).toDF("id", "embedding")
})

问题分析：由于wordToVectorMap是一个driver端数据量较大的本地map，idToWordVector在多个executor的计算过程中，涉及到序列化和反序列化的问题，与闭包的概念相关，idToWordVector需要将本地的wordToVectorMap进行序列化，在此过程中，超过了字节缓冲器的最大限制，从而发生OOM。

3. 解决方案

// 代码2
val sql = s"select id, text from table where day = '2020-08-20'"

val wordToVector: RDD[(String, Array[Float])] = ...
val EMBEDDING_SIZE = 300
val zeroVector = new Array[Float](EMBEDDING_SIZE)
val wordToId: RDD[(String, String)] = hiveContext.sql(sql).rdd.flatMap(x => {

    val id = x.getString(0)
    val wordList: Array[String] = x.getString(1).split(" ")
    
    val listBuffer = new ListBuffer[(String, String)]()
    for (word <- wordList) {
        listBuffer.append((word, id))
    }
    listBuffer
})

val idToWordVectorRdd: RDD[(String, Array[Float])] = wordToId.join(wordToVector).map(x => (x._2._1, x._2._2))
    .reduceByKey((a, b) => {
        val result = new Array[Float](a.length)
        for (i <- a.indices) {
            result(i) = a(i) + b(i)
        }
        result
    })

val idToWordVector: DataFrame= idToWordVectorRdd
    .map(x => (x._1, x._2.mkString(",")))
    .toDF("id", "embedding")

将代码1中的wordToVectorMap表示为rdd形式wordToVector。同时创建一个rdd wordToId表示word到id的映射关系。将wordToId和wordToVector进行关联，得到<id, vector>的rdd表示，然后进行聚合操作即可。

4. 总结

这种解决方案的本质思想即人工进行了map和reduce的操作，提高了数据的分块和执行的并行度。