Spark中高效计算top-k元素

我有一个类似于以下内容的数据框：

+---+-----+-----+
|key|thing|value|
+---+-----+-----+
| u1|  foo|    1|
| u1|  foo|    2|
| u1|  bar|   10|
| u2|  foo|   10|
| u2|  foo|    2|
| u2|  bar|   10|
+---+-----+-----+

并希望得到以下结果：

+---+-----+---------+----+
|key|thing|sum_value|rank|
+---+-----+---------+----+
| u1|  bar|       10|   1|
| u1|  foo|        3|   2|
| u2|  foo|       12|   1|
| u2|  bar|       10|   2|
+---+-----+---------+----+

目前，有类似的代码：

val df = Seq(("u1", "foo", 1), ("u1", "foo", 2), ("u1", "bar", 10), ("u2", "foo", 10), ("u2", "foo", 2), ("u2", "bar", 10)).toDF("key", "thing", "value")

 // calculate sums per key and thing
 val aggregated = df.groupBy("key", "thing").agg(sum("value").alias("sum_value"))

 // get topk items per key
 val k = lit(10)
 val topk = aggregated.withColumn("rank", rank over  Window.partitionBy("key").orderBy(desc("sum_value"))).filter('rank < k)

不过这段代码很效率低下。窗口函数生成一个总订单的项目并导致巨大的洗牌.

如何更有效地计算 top-k 项？ 也许使用近似函数，即类似于草图https://datasketches.github.io/ https://datasketches.github.io/ or https://spark.apache.org/docs/latest/ml-frequent-pattern-mining.html https://spark.apache.org/docs/latest/ml-frequent-pattern-mining.html

这是推荐系统的经典算法。

case class Rating(thing: String, value: Int) extends Ordered[Rating] {
  def compare(that: Rating): Int = -this.value.compare(that.value)
}

case class Recommendation(key: Int, ratings: Seq[Rating]) {
  def keep(n: Int) = this.copy(ratings = ratings.sorted.take(n))
}

val TOPK = 10

df.groupBy('key)
  .agg(collect_list(struct('thing, 'value)) as "ratings")
  .as[Recommendation]
  .map(_.keep(TOPK))

您还可以在以下位置查看源代码：

• Spotify 大数据 Rosetta 代码 /TopItemsPerUser.scala https://github.com/spotify/big-data-rosetta-code/blob/master/src/main/scala/com/spotify/bdrc/pipeline/TopItemsPerUser.scala，这里有几个针对 Spark 或 Scio 的解决方案
• Spark MLLib /TopByKeyAggregator.scala https://github.com/apache/spark/blob/master/mllib/src/test/scala/org/apache/spark/ml/recommendation/TopByKeyAggregatorSuite.scala，被认为是使用他们的推荐算法时的最佳实践，看起来他们的例子仍然使用RDD though.

import org.apache.spark.mllib.rdd.MLPairRDDFunctions._

sc.parallelize(Array(("u1", ("foo", 1)), ("u1", ("foo", 2)), ("u1", ("bar", 10)), ("u2", ("foo", 10)),
  ("u2", ("foo", 2)), ("u2", ("bar", 10))))
  .topByKey(10)(Ordering.by(_._2))