Scala 有效地实现spark的takeByKey

我有一个类型为

RDD[（k:Int，v:String）]

的RDD。我想对每个键

k

使用多达1000个元组，这样我就有了

[（k，v）]

，其中没有键出现超过1000次有没有一种方法可以避免先调用groupBy的性能损失？我想不出一个好的方法来聚合这些值，以避免一个完整的groupBy导致我的工作失败

天真的方法：

def takeByKey(rdd: RDD[(K,V)], n: Int) : RDD[(K,V)] = {
    rdd.groupBy(_._1).mapValues(_.take(n)).flatMap(_._2)
}

我正在寻找一种更有效的方法来避免分组：

takeByKey(rdd: RDD[(K,V)], n: Int) : RDD[(K,V)] = {
    //use reduceByKey, foldByKey, etc..??
}

---

这是迄今为止我开发的最好的解决方案，但它不进行类型检查

def takeByKey(rdd: RDD[(K,V)], n: Int) : RDD[(K,V)] = {
      rdd.foldByKey(List[V](), ((acc, elem) => if (acc.length >= n) acc else elem._2 :: acc)).flatMap(t => t._2.map(v => (t._1, v)))
}

编辑。 我提出了一个稍微好一点的解决方案，似乎正在发挥作用：

takeByKey(rdd: RDD[(K,V)], n: Int) : RDD[(K,V)] = {
    rdd.mapValues(List(_))
       .reduceByKey((x,y) => if(x.length >= n) x 
                             else if(y.length >= n) y 
                             else (x ++ y).take(n))
       .flatMap(t => t._2.map(v => (t._1, v)))
}

---

这是迄今为止我提出的最好的解决方案

takeByKey(rdd: RDD[(K,V)], n: Int) : RDD[(K,V)] = {
    rdd.mapValues(List(_))
       .reduceByKey((x,y) => if(x.length >= n) x 
                             else if(y.length >= n) y 
                             else (x ++ y).take(n))
       .flatMap(t => t._2.map(v => (t._1, v)))
}

---

它不会像groupByKey方法那样耗尽内存并死掉，但它仍然很慢。

您当前的解决方案是朝着正确方向迈出的一步，但它仍然非常低效，原因至少有三个：

• mapValues（List（））

  创建大量临时

  List

  对象
• 线性

  序列的
  长度
  类似
  列表
  为O（N）

• x++y

  再次创建大量临时对象

最简单的方法是用可变缓冲区替换

List

，该缓冲区具有恒定的时间

长度

。一个可能的例子如下：

import scala.collection.mutable.ArrayBuffer

rdd.aggregateByKey(ArrayBuffer[Int]())(
  (acc, x) => if (acc.length >= n) acc else acc += x,
  (acc1, acc2) => {
    val (xs, ys) = if (acc1.length > acc2.length) (acc1, acc2) else (acc2, acc1)
    val toTake = Math.min(n - xs.length, ys.length)
    for (i <- 0 until toTake) {
      xs += ys(i)
    }
    xs         
  }
)

与

它不会影响性能，但稍微干净一些

.flatMap(t => t._2.map(v => (t._1, v)))

.flatMapValues(x => x)  // identity[Seq[V]]