在Spark数据集中滚动自己的reduceByKey

时间:2016-07-14 19:56:49

标签: scala apache-spark mapreduce

除了RDD之外,我还试图学习更多使用DataFrame和DataSet。对于RDD,我知道我可以someRDD.reduceByKey((x,y) => x + y),但我没有看到数据集的功能。所以我决定写一个。

someRdd.map(x => ((x.fromId,x.toId),1)).map(x => collection.mutable.Map(x)).reduce((x,y) => {
  val result = mutable.HashMap.empty[(Long,Long),Int]
  val keys = mutable.HashSet.empty[(Long,Long)]
  y.keys.foreach(z => keys += z)
  x.keys.foreach(z => keys += z)
  for (elem <- keys) {
    val s1 = if(x.contains(elem)) x(elem) else 0
    val s2 = if(y.contains(elem)) y(elem) else 0
    result(elem) = s1 + s2
  }
  result
})

然而,这会将所有内容返回给驱动程序。你怎么写这个来回归Dataset?也许mapPartition并在那里做?

注意这个编译但是没有运行,因为它没有Map的编码器

2 个答案:

答案 0 :(得分:26)

我认为你的目标是将这个习语翻译成数据集:

rdd.map(x => (x.someKey, x.someField))
   .reduceByKey(_ + _)

// => returning an RDD of (KeyType, FieldType)

目前,我在数据集API中找到的最接近的解决方案如下所示:

ds.map(x => (x.someKey, x.someField))          // [1]
  .groupByKey(_._1)                            
  .reduceGroups((a, b) => (a._1, a._2 + b._2))
  .map(_._2)                                   // [2]

// => returning a Dataset of (KeyType, FieldType)

// Comments:
// [1] As far as I can see, having a map before groupByKey is required
//     to end up with the proper type in reduceGroups. After all, we do
//     not want to reduce over the original type, but the FieldType.
// [2] required since reduceGroups converts back to Dataset[(K, V)]
//     not knowing that our V's are already key-value pairs.

看起来不是很优雅,根据快速基准,它的表现也差得多,所以也许我们在这里遗漏了一些......

注意:替代方案可能是使用groupByKey(_.someKey)作为第一步。问题是使用groupByKey会将类型从常规Dataset更改为KeyValueGroupedDataset。后者没有常规map功能。相反,它提供mapGroups,这似乎不太方便,因为它将值包装到Iterator并根据文档字符串执行shuffle。

答案 1 :(得分:7)

更有效的解决方案在mapPartitions之前使用groupByKey来减少重排量(请注意,这与reduceByKey不完全相同,但我认为传递更为灵活功能要求数据集由元组组成)。

def reduceByKey[V: ClassTag, K](ds: Dataset[V], f: V => K, g: (V, V) => V)
  (implicit encK: Encoder[K], encV: Encoder[V]): Dataset[(K, V)] = {
  def h[V: ClassTag, K](f: V => K, g: (V, V) => V, iter: Iterator[V]): Iterator[V] = {
    iter.toArray.groupBy(f).mapValues(_.reduce(g)).map(_._2).toIterator
  }
  ds.mapPartitions(h(f, g, _))
    .groupByKey(f)(encK)
    .reduceGroups(g)
}

根据数据的形状/大小,这比reduceByKey的效果在1秒内,与2x一样快groupByKey(_._1).reduceGroups。仍有改进的余地,欢迎提出建议。