Spark TopK 问题解决-使用最小堆

参考资料：
《Spark 大数据处理》 by 高彦杰

整个排序取 TopK 的实现：

object TopK0 {

  val K = 3

  def main(args: Array[String]) {
    // 执行 wordcount
    val conf = new SparkConf().setAppName("TopK0")
    val spark = new SparkContext(conf)
    val textRDD = spark.textFile("hdfs://10.0.8.162:9000/home/yuzx/input/wordcount.txt")
    textRDD.flatMap(line => line.split(" "))
      .map(word => (word, 1))
      .reduceByKey(_ + _)
      .map { pair => pair.swap }
      .sortByKey(true, 2)
      .top(3)
      .foreach(println)

    spark.stop()

  }

}

基于最小堆的实现：

最大/小堆，对应的数据结构优先级队列，PriorityQueue，不光 Java 中有，Scala 中也有，当然 c++ 中也有

object TopK {

  val K = 3

  val ord = Ordering.by[(String, Int), Int](_._2).reverse

  def main(args: Array[String]) {
    // 执行 wordcount
    val conf = new SparkConf().setAppName("TopK")
    val spark = new SparkContext(conf)
    val textRDD = spark.textFile("hdfs://10.0.8.162:9000/home/yuzx/input/wordcount.txt")
    val countRes = textRDD.flatMap(line => line.split(" ")).map(word => (word, 1)).reduceByKey(_ + _)

    // debug mapreduce 的结果
    countRes.foreach(println)

    /*
     每个 RDD 分区内进行 TOP K 计算
     需要每个分区内有自己的桶，如果整个程序使用一个 heap（将 heap 设定为成员变量） 会不正确
     为什么呢？
      */
    val topk = countRes.mapPartitions(iter => {
      val heap = new mutable.PriorityQueue[(String, Int)]()(ord)
      while (iter.hasNext) {
        val n = iter.next
        println("分区计算：" + n)
        putToHeap(heap, n)
      }

      heap.iterator
    }).collect()

    println("分区结果：")
    topk.foreach(println)

    // 每个分区的 TOP K 合并，计算总的 TopK
    val heap = new mutable.PriorityQueue[(String, Int)]()(ord)
    val iter = topk.iterator
    while (iter.hasNext) {
      putToHeap(heap, iter.next)
    }

    println("最终结果：")
    while (heap.nonEmpty) {
      println(heap.dequeue())
    }

    spark.stop()

  }

  def putToHeap(heap: mutable.PriorityQueue[(String, Int)], iter: (String, Int)): Unit = {
    if (heap.nonEmpty && heap.size >= K) {
      if (heap.head._2 < iter._2) {
        heap += iter
        heap.dequeue()
      }
    } else {
      heap += iter
    }
  }

}

有什么不对的地方，欢迎指出～～