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的操作,提高了数据的分块和执行的并行度。