Spark算子之foldByKey

在学习foldByKey这个算子的时候，发现网上好多文章的内容相互冲突，于是决定自己实践一边，以理解这个算子是怎么运行的。

foldByKey

def foldByKey(zeroValue: V, numPartitions: Int)(func: (V, V) => V): RDD[(K, V)] = self.withScope {
    foldByKey(zeroValue, new HashPartitioner(numPartitions))(func)
  }

def foldByKey(zeroValue: V)(func: (V, V) => V): RDD[(K, V)] = self.withScope {
    foldByKey(zeroValue, defaultPartitioner(self))(func)
  }

这是源码中对foldByKey的定义，主要是zeroValue这个点容易造成误解：

我们先初始化一个RDD

scala> var rdd = sc.makeRDD(Array(("A",1),("A",2),("B",1),("B",2),("C",1)))
rdd: org.apache.spark.rdd.RDD[(String, Int)] = ParallelCollectionRDD[6] at makeRDD at <console>:24

查看这个RDD的分区及分区中的元素

scala> rdd.mapPartitionsWithIndex((index, iter) => {
     |       var rddmap = scala.collection.mutable.Map[String, List[(String, Int)]]()
     |       while (iter.hasNext) {
     |         var elem = iter.next()
     |         var partNum = index + "_"
     |         if (rddmap.contains(partNum)) {
     |           var elems = rddmap(partNum)
     |           elems ::= elem
     |           rddmap(partNum) = elems
     |         } else {
     |           var newlist = List[(String,Int)]()
     |           newlist ::= elem
     |           rddmap(partNum) = newlist
     |         }
     |       }
     |       rddmap.toIterator
     |     }).collect()
res11: Array[(String, List[(String, Int)])] = Array((0_,List((A,2), (A,1))), (1_,List((C,1), (B,2), (B,1))))

可以看到，目前是有两个分区：
第一个分区中包含(A,2), (A,1)；
第二个分区中包含(C,1), (B,2), (B,1)

这时候我们用foldByKey算子，zeroValue设为2

scala>  rdd.foldByKey(2) { (x1, x2) => x1 + x2 }.collect()
res12: Array[(String, Int)] = Array((B,5), (A,5), (C,3))

再来看看另一种情况

在初始化RDD的过程中，将分区数设为3

scala> var rdd = sc.makeRDD(Array(("A", 1), ("A", 2), ("B", 1), ("B", 2), ("C", 1)),3)
rdd: org.apache.spark.rdd.RDD[(String, Int)] = ParallelCollectionRDD[11] at makeRDD at <console>:24

再查看下每个分区中的元素

scala> rdd.mapPartitionsWithIndex((index, iter) => {
     |       var rddmap = scala.collection.mutable.Map[String, List[(String, Int)]]()
     |       while (iter.hasNext) {
     |         var elem = iter.next()
     |         var partNum = index + "_"
     |         if (rddmap.contains(partNum)) {
     |           var elems = rddmap(partNum)
     |           elems ::= elem
     |           rddmap(partNum) = elems
     |         } else {
     |           var newlist = List[(String,Int)]()
     |           newlist ::= elem
     |           rddmap(partNum) = newlist
     |         }
     |       }
     |       rddmap.toIterator
     |     }).collect()
res13: Array[(String, List[(String, Int)])] = Array((0_,List((A,1))), (1_,List((B,1), (A,2))), (2_,List((C,1), (B,2))))

第一个分区中只有 （A，1）
第二个分区中有 （B，1）（A，2）
第三个分区中有 （C，1）（B，2）
这时我们再使用foldByKey算子，zeroValue还是设为2

scala> rdd.foldByKey(2) { (x1, x2) => x1 + x2 }.collect()
res14: Array[(String, Int)] = Array((B,7), (C,3), (A,7))

出现了和第一次不一样的结果

由此我们已经可以推断出foldByKey的机制了，我的理解是：
取RDD的每一个分片，在每一个分片中，先根据你定义的映射，用zeroValue对不同key对应的value做一次初始化，再对剩下的value值做映射。

在第一次操作中，第一个分区中包含(A,2), (A,1)，先做初始化 2+2 = 4 ，再对剩下的值做累加， 最后得到 （A，4+1 = 5）

在第二次操作中，第一个分区中只有 （A，1），初始化为（A，2+1 = 3），第二个分区中有 （B，1）（A，2），因为是不同分区，所以也对A进行初始化（A，2+2 = 4），最后两个分区的结果统计到一起就是（A，4+3 = 7），BC同理。