版权归本人及公司所有,转载请注明出处。
https://blog.csdn.net/aicodex/article/details/90297609
背景
最近在做推荐系统,其中要对数据做召回,需要快速为每个待推荐的“用户”召回100个左右的“电影”。此时一个简单的想法就是计算每个“用户”和“电影”的“相关度”,取top100。假设“用户”和“电影”都有标签。例如“科幻”,“动作”。那么在计算top100,可以只计算同key的相关度。
实现的时候很简单,假设你把用户数据读取好了,电影数据也读取好了
val user = sc.parallelize(Array(("科幻","小明"),("科幻","小红"),("剧情","小王"),("剧情","小李"),("剧情","小张")))
val movie = sc.parallelize(Array(("科幻","复联4"),("剧情","功夫"),("剧情","弱点")))
很自然想到的方法就是把他们按key求交。分别计算“用户”和“电影”的相关度。
val relevant = user.join(movie).map(e=>computeScore(e._2._1,e._2._2))
但是在这时你忽然发现喜欢剧情类电影的用户如此之多、剧情类的电影也如此多。最后导致了严重的数据倾斜。就像上述数据,join之后得到了。
val joinResult = user.join(movie).collect
joinResult: = Array((剧情,(小李,功夫)), (剧情,(小李,弱点)), (剧情,(小王,功夫)), (剧情,(小王,弱点)), (剧情,(小张,功夫)), (剧情,(小张,弱点)), (科幻,(小明,复联4)), (科幻,(小红,复联4)))
没错,join操作其实就是“按key做笛卡尔积”。看一下spark的join的实现就明白了,join操作本质上是做了两步cogroup和flatMap:
//spark的join源码
def join[W](other: RDD[(K, W)], partitioner: Partitioner): RDD[(K, (V, W))] = self.withScope {
this.cogroup(other, partitioner).flatMapValues( pair =>
for (v <- pair._1.iterator; w <- pair._2.iterator) yield (v, w)
)
}
cogroup和groupByKey类似,就是cogroup的两个rdd按key做group操作,得到一个(key,valueArr)的元祖。这个flatMapValues就是按value铺平。key复制原来的。此时假如科幻电影是1000部,喜欢科幻的用户1000个。那么你会得到100w个pair。倘若剧情电影1w部,喜欢剧情电影的人1w个。那么你会得到1亿个pair。想想一下这些pair会出现在一个节点上(cogroup的时候同key一定会被hash到一个节点)。而且hash到哪个节点完全是随机的!!假设你的节点还挺强,1亿个还撑不爆你的节点,随机hash两个1亿全跑一个节点去了。就直接无法执行了,即使可以,也会因为结果太大spill到硬盘巨慢无比。
解决思路一
逻辑源于生活,高于生活。倘若给你一堆石头,有大的,有小的。让你尽可能平均的放到背包里。假如有两个石头非常大,你是否就应该把它尽可能别放到一个背包里,防止分配严重不均。这个问题是一个NP-Hard问题。一般称为K-Partition问题。自然而然想到的方法和维基上这个贪心算法很接近。我们把石头按重到轻依次放入目前最轻的背包。尽可能让他平均来解决数据倾斜。(绝对平均的代价太高了,NP-Hard问题意味着你不可能在多项式时间求解,只能暴力破解)
我们这个实际问题比这个稍微复杂,因为在最终计算的过程中,得到的元素个数是“用户”数 x “电影”数。因此我们需要定义一个称重函数,使用者最了解自己的数据如何“称重”让他自己去称,我们要做的是按照“重量”,尽可能均分到每个节点上。话不多说直接上代码:
import org.apache.spark.Partitioner
import org.apache.spark.rdd.RDD
import scala.reflect.ClassTag
/*
* 自定义解决数据不平衡的分区器。通常数据不平衡经常出现在两个RDD求join、或一个RDD groupByKey的过程中。
* 而HashPartitioner是根据Key的hashCode对分区数取余,可以看做是随机分区。最后数据是按key均匀分布的。
* 但是每个key对应的value不一定均衡。尤其是join的时候,如果两个key对应的数据最后还要求笛卡尔积铺平的话,
* key对应的元素个数可能就是rdd1的key对应元素个数*rdd2的key对应的元素个数。也可能计算量是指数级甚至乘方级(例如求两两相似度,迭代梯度下降)。
* 因此需要根据"计算量"去尽可能平均划分分区,才是最优策略。当然在这个基础上,还是需要保证同key的分区是一个。
* rdd1是你输入的第一个rdd,valueWeightFunc1是对他key对应的值进行评估的函数(这个评估函数必须很快,否则还不如不用这个分区方式)
* rdd2是你输入的第二个rdd,valueWeightFunc2是对他key对应的值进行评估的函数
* mergeWeightFunc,是输入rdd1的key对应的weight和rdd2的key对应的weight。输出最终weight的函数。
* 该分区器会根据key对应的最终weight去分区,让每个分区的weight尽可能平均。
* 平均分配weight的任务像是一个经典问题:Partition problem.
* @see <a href="https://en.wikipedia.org/wiki/Partition_problem">https://en.wikipedia.org/wiki/Partition_problem</a>
* Partition problem是一个NP-Hard的问题,本代码的实现类似于该维基百科的近似算法中的贪心算法。使用了一个小顶堆来实现快速查找最小总和的子集。
* 这里主要是考虑两个key相互影响的情况。rdd1,rdd2 key不相互重叠的部分,会采用HashPartitioner的方法。
*
*/
class KPartitioner[K, V, W](rdd1: RDD[(K, V)], rdd2: RDD[(K, W)], valueWeightFunc1: V => Int, valueWeightFunc2: W => Int, mergeWeightFunc: (Int, Int) => Int, partitions: Int)(implicit kt: ClassTag[K], vt: ClassTag[V], wt: ClassTag[W]) extends Partitioner {
val keyPartMap = {
val rdd1KeyWeight = rdd1.mapValues(valueWeightFunc1).reduceByKey(_ + _)
val rdd2KeyWeight = rdd2.mapValues(valueWeightFunc2).reduceByKey(_ + _)
val mergeKeyWeight = rdd1KeyWeight.join(rdd2KeyWeight).mapValues(e => mergeWeightFunc(e._1, e._2))
val keyCountArr = mergeKeyWeight.sortBy(_._2, false).collect()
val minHeap = new MinHeap(keyCountArr.take(partitions).map(e => (List(e._1), e._2.toLong)))
keyCountArr.slice(partitions, keyCountArr.length).foreach(e => minHeap.addKey(e._1, e._2))
minHeap.data.map(_._1).zipWithIndex.flatMap(e => e._1.map(f => (f, e._2))).toMap
}
override def numPartitions: Int = partitions
override def getPartition(key: Any): Int = keyPartMap.getOrElse(key.asInstanceOf[K], nonNegativeMod(key.hashCode, numPartitions))
def nonNegativeMod(x: Int, mod: Int): Int = {
val rawMod = x % mod
rawMod + (if (rawMod < 0) mod else 0)
}
class MinHeap(var data: Array[(List[K], Long)]) {
buildHeap()
private def buildHeap(): Unit = {
var i = data.length / 2 - 1
while (i >= 0) {
heapify(i)
i -= 1
}
}
private def heapify(i: Int): Unit = {
val l = left(i)
val r = right(i)
var smallest = i
if (l < data.length && data(l)._2 < data(i)._2) smallest = l
if (r < data.length && data(r)._2 < data(i)._2) smallest = r
if (i == smallest) return
swap(i, smallest)
heapify(smallest)
}
private def right(i: Int) = (i + 1) << 1
private def left(i: Int) = ((i + 1) << 1) - 1
private def swap(i: Int, j: Int): Unit = {
val tmp = data(i)
data(i) = data(j)
data(j) = tmp
}
def getRoot = data(0)
def setRoot(root: (List[K], Long)): Unit = {
data(0) = root
heapify(0)
}
def addKey(key: K, size: Long): Unit = {
val oldRoot = getRoot
setRoot((oldRoot._1 :+ key, oldRoot._2 + size))
}
}
}
object KPartitioner {
def apply[K, V, W](rdd1: RDD[(K, V)], rdd2: RDD[(K, W)], valueWeightFunc1: V => Int, valueWeightFunc2: W => Int, mergeWeightFunc: (Int, Int) => Int, partitions: Int)(implicit kt: ClassTag[K], vt: ClassTag[V], wt: ClassTag[W]) = new KPartitioner[K, V, W](rdd1, rdd2, valueWeightFunc1, valueWeightFunc2, mergeWeightFunc, partitions)
def apply[K, V](rdd: RDD[(K, V)], valueWeightFunc: V => Int, partitions: Int)(implicit kt: ClassTag[K], vt: ClassTag[V]) = new KPartitioner[K, V, V](rdd, rdd, valueWeightFunc, f => 0, (x, y) => x + y, partitions)
}
//测试代码
import xxx.KPartitioner
import org.apache.spark.{SparkConf, SparkContext}
import org.junit.Test
@Test
class KPartitionerTest {
@Test
def getPartitionTest(): Unit = {
val conf = new SparkConf().setAppName("mySpark")
conf.setMaster("local")
val sc = new SparkContext(conf)
val inputList = List(("a", 4), ("b", 1), ("c", 3), ("d", 2), ("e", 7), ("f", 10), ("g", 20), ("h", 4), ("i", 5))
val inputMap = inputList.toMap
val testRDD = sc.parallelize(inputList)
val part = KPartitioner[String, Int](testRDD, f => f, 4)
//看每个分区分配的元素,以及合理性。期待20单独一个分区、10 2一个、7 4 1一个、5 4 3一个。
part.keyPartMap.toList.map(_.swap).groupBy(_._1).mapValues(_.map(e => (e._2, inputMap.getOrElse(e._2, -1)))).foreach(println)
}
}
代码中的小顶堆只是为了快速找到“最轻背包”的方法。代码中也附了简单的测试代码。只需要和hashPartitioner类似使用就行。join/cogroup/groupByKey的时候都可以用。最终效果是各个分区元素都尽可能平均了,但是计算这个较优分区,多付出了“称重”的时间代价。这个方法有一个缺陷就是,倘若有一个石头还是太大,你哪怕把它单独放在一个背包,他也过分倾斜,那么是没法用的。
解决思路二
还是我遇到的这个问题,我们注意到这样一件事:join是由cogroup+类似flatMap实现的。假如1w个喜欢剧情电影的人和1w部剧情电影。cogroup到一起也不过区区2w的数据量,这么点量很难导致数据不平衡。导致数据不平衡的根本原因是我们需要对他进行两两相似度计算。所以flatMap这一步操作又是几乎是必须的(即使你不flat也得计算1亿次相关度)。但是反过来,1亿条数据对于spark来说真不多,多的原因是把它分配到一个节点上去了。所以只要不让他在一个节点上就行。
说的轻巧,不在一个节点,我怎么保证一个key的数据还在一起呢?如果我在节点1上5000个人5000个电影,我如何让他知道节点2上的5000人和5000部电影呢。他们被隔断了,广播代价太高。其实思路说难也难,说简单也简单。举个简单的例子:
* [a,b,c,d] · [1,2] -> [(a,1)(a,2)(b,1)(b,2)(c,1)(c,2)(d,1)(e,2)]
* 直接铺平有8个字母,8个数字。而先按照左侧铺平,和右侧铺平分别为:
* (a,[1,2]) ([a,b,c,d],1)
* (c,[1,2]) ([a,b,c,d],2)
* (d,[1,2])
* (e,[1,2])
* 先铺平左侧。那么字母还是4个。数字变成8个。先铺平右侧字母变成8个数字是2个。
所以先铺平一侧,再把它利用repatition随机分发到各个节点上,就可以大大减小每个节点的数据量。当然,如果铺平一侧还是过大,我可以把铺平的一侧进行拆分再多铺平一次。话不多说直接上代码:
import org.apache.spark.rdd.RDD
import scala.reflect.ClassTag
object RDDFunctionExtend {
implicit class PairRDDCovertor[K, V](val rdd: RDD[(K, V)]) extends AnyVal {
def cartesianByKeyWith[W](other: RDD[(K, W)], splitNumLeft: Int)(implicit kt: ClassTag[K], vt: ClassTag[V], wt: ClassTag[W]) = {
rdd.cogroup(other)
.flatMapValues(pair => pair._1.grouped(splitNumLeft).map(f => (f, pair._2)))
.repartition(rdd.sparkContext.defaultParallelism)
.flatMapValues(pair => for (v <- pair._1.iterator; w <- pair._2.iterator) yield (v, w))
}
def cartesianByKeyWith[W](other: RDD[(K, W)], splitNumLeft: Int, splitNumRight: Int)(implicit kt: ClassTag[K], vt: ClassTag[V], wt: ClassTag[W]) = {
rdd.cogroup(other)
.flatMapValues(pair => for (v <- pair._1.grouped(splitNumLeft); w <- pair._2.grouped(splitNumRight)) yield (v, w))
.repartition(rdd.sparkContext.defaultParallelism)
.flatMapValues(pair => for (v <- pair._1.iterator; w <- pair._2.iterator) yield (v, w))
}
def cartesianByKeyWith[W](other: RDD[(K, W)])(implicit kt: ClassTag[K], vt: ClassTag[V], wt: ClassTag[W]) = {
rdd.cogroup(other)
.flatMapValues(e => e._1.map(f => (f, e._2)))
.repartition(rdd.sparkContext.defaultParallelism)
.flatMapValues(e => e._2.map(f => (e._1, f)))
}
}
}
我定义了这么一个“隐式转换”能吧pairRDD转化成我定义的类,用的时候只需要import这个隐式类即可。这样给pairRDD拓展了一个cartesianByKeyWith接口,可以让他铺平一侧,或者是拆分了再铺平。来替代join。如果需要leftOuterJoin或者rightOuterJoin或是fullOuterJoin的,可以自己照葫芦画瓢实现一个。