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

来源：互联网发布：八爪鱼采集软件免费版编辑：程序博客网时间：2024/06/06 18:34

参考资料：
《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    }  }}

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

0 0