Spark Datasetで独自のreduceByKeyをローリング

Question

私はRDDに加えてDataFramesとDataSetsの使用方法を学びたいと考えています。 RDDの場合、私はsomeRDD.reduceByKey((x,y) => x + y)を行うことができますが、私はデータセットのためのその機能は表示されません。だから私は1つ書くことにしました。

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まだ

Answer 1

ためのエンコーダを持っていないため

注これはコンパイルが、実行されません。私はあなたの目標は、データセットには、このイディオムを翻訳することであると仮定します。

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

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

現在、最も近いソリューションIデータセット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にラップして、docstringに従ってシャッフルを実行するので非常に便利ではありません。

Answer 2

より効率的なソリューションは、シャッフルの量を減らすためにgroupByKey前mapPartitionsを使用する（これは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) 
} 

あなたのデータの形状/サイズに応じて、これは約2x早くgroupByKey(_._1).reduceGroupsとしてreduceByKeyの性能を1秒以内である、と。まだ改善の余地があるので、提案は大歓迎です。