web-dev-qa-db-fra.com

Déployer votre propre ReduceByKey dans Spark Dataset

J'essaie d'apprendre à utiliser davantage les DataFrames et DataSets en plus des RDD. Pour un RDD, je sais que je peux faire someRDD.reduceByKey((x,y) => x + y), mais je ne vois pas cette fonction pour Dataset. J'ai donc décidé d'en écrire un.

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
})

Cependant, cela renvoie tout au pilote. Comment écririez-vous ceci pour retourner un Dataset? Peut-être mapPartition et le faire là-bas?

Notez que cela se compile mais ne s'exécute pas car il n'a pas encore d'encodeur pour Map

scalaapache-sparkmapreduce
21
Carlos Bribiescas

Je suppose que votre objectif est de traduire cet idiome en ensembles de données:

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

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

Actuellement, la solution la plus proche que j'ai trouvée avec l'API Dataset ressemble à ceci:

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.

Ne semble pas très élégant et selon une référence rapide, il est également beaucoup moins performant, alors peut-être qu'il nous manque quelque chose ici ...

Remarque: Une alternative pourrait être d'utiliser groupByKey(_.someKey) comme première étape. Le problème est que l'utilisation de groupByKey change le type d'un Dataset normal en KeyValueGroupedDataset . Ce dernier n'a pas de fonction régulière map. Au lieu de cela, il propose un mapGroups, ce qui ne semble pas très pratique car il encapsule les valeurs dans un Iterator et effectue un shuffle en fonction de la docstring.

28
bluenote10

Une solution plus efficace utilise mapPartitions avant groupByKey pour réduire la quantité de brassage (notez que ce n'est pas exactement la même signature que reduceByKey mais je pense qu'il est plus flexible de passer un fonction qui nécessite que l'ensemble de données se compose d'un Tuple).

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)
}

Selon la forme/taille de vos données, cela se situe à moins d'une seconde des performances de reduceByKey, et environ 2x Aussi rapidement qu'une groupByKey(_._1).reduceGroups. Il y a encore de la place pour des améliorations, donc des suggestions seraient les bienvenues.

7
Justin Raymond