foldByKey, aggregateByKey et combineByKey en Spark

Calculer une somme par clé avec reduceByKey est immédiat. Calculer une moyenne par clé ne l’est pas — et la raison est structurelle : reduceByKey exige que le résultat soit du même type que les valeurs d’entrée. Pour une moyenne, il faut transporter deux informations (une somme et un compte), donc changer de type en cours de route.

C’est exactement le problème que résolvent foldByKey, aggregateByKey et combineByKey. Trois outils, une même famille, un ordre croissant de puissance et de verbosité. On les compare sur le même exemple.

Le jeu de données

Des notes par étudiant :

val notes = sc.parallelize(List(
  ("Ali", 15), ("Sara", 12), ("Ali", 9),
  ("Yassine", 18), ("Sara", 16), ("Ali", 12)
))

Objectif : la moyenne par prénom. Résultat attendu :

Ali     : 12.0   // (15 + 9 + 12) / 3
Sara    : 14.0   // (12 + 16) / 2
Yassine : 18.0   // 18 / 1

Pourquoi reduceByKey ne peut pas

notes.reduceByKey(_ + _)
// ("Ali", 36), ("Sara", 28), ("Yassine", 18)

La somme, oui. Mais le compte a disparu — il aurait fallu le porter en même temps. La signature l’interdit :

def reduceByKey(func: (V, V) => V): RDD[(K, V)]

Entrée V, sortie V. Même type. Pour une moyenne, il faut passer de Int à (Int, Int) — la famille aggregateByKey existe pour ça.

La famille, en une table

Toutes les quatre effectuent le combine local avant le shuffle — c’est la propriété qui les distingue de groupByKey et qui fait leur intérêt en performance.

foldByKey : reduceByKey avec un élément neutre

notes.foldByKey(0)(_ + _).collect()
// ("Ali", 36), ("Sara", 28), ("Yassine", 18)

La valeur 0 est l’élément neutre de l’opération, appliqué au démarrage de chaque partition. L’intérêt sur reduceByKey est étroit mais réel :

// Le maximum, avec une borne inférieure explicite
notes.foldByKey(Int.MinValue)(math.max).collect()

// Une multiplication, dont le neutre est 1 et non 0
scores.foldByKey(1)(_ * _).collect()

Attention : la valeur initiale est appliquée par partition, pas une fois globalement. Avec foldByKey(10)(_ + _) sur trois partitions contenant la même clé, vous ajoutez 30, pas 10. C’est pour cette raison que l’élément doit être neutre pour l’opération — sinon le résultat dépend du nombre de partitions.

aggregateByKey : changer de type, enfin

C’est l’outil de la moyenne. Sa signature demande deux fonctions :

def aggregateByKey[U](valeurInitiale: U)(
  dansLaPartition: (U, V) => U,      // accumule une valeur dans l'accumulateur
  entrePartitions: (U, U) => U       // fusionne deux accumulateurs
): RDD[(K, U)]

L’accumulateur U peut être de n’importe quel type. Pour une moyenne, c’est un couple (somme, compte) :

val moyennes = notes
  .aggregateByKey((0, 0))(
    // dans la partition : on ajoute une note à l'accumulateur
    (acc, note) => (acc._1 + note, acc._2 + 1),
    // entre partitions : on additionne deux accumulateurs
    (a, b) => (a._1 + b._1, a._2 + b._2)
  )
  .mapValues { case (somme, compte) => somme.toDouble / compte }

moyennes.collect().foreach(println)
// (Ali, 12.0)
// (Sara, 14.0)
// (Yassine, 18.0)

Les deux fonctions ont des rôles bien distincts, et c’est le point qui bloque tout le monde au début :

dansLaPartition mélange un accumulateur et une valeur brute. entrePartitions fusionne deux accumulateurs. Les types diffèrent, d’où les deux fonctions — et c’est ce qui permet de ne transporter que (somme, compte) sur le réseau au lieu de toutes les notes.

// Collecter les valeurs distinctes par clé, sans groupByKey
notes.aggregateByKey(Set.empty[Int])(_ + _, _ ++ _)

// Min et max simultanément
notes.aggregateByKey((Int.MaxValue, Int.MinValue))(
  (acc, v) => (math.min(acc._1, v), math.max(acc._2, v)),
  (a, b)   => (math.min(a._1, b._1), math.max(a._2, b._2))
)

// Statistiques complètes : compte, somme, somme des carrés → moyenne et écart-type
notes.aggregateByKey((0L, 0.0, 0.0))(
  (acc, v) => (acc._1 + 1, acc._2 + v, acc._3 + v * v),
  (a, b)   => (a._1 + b._1, a._2 + b._2, a._3 + b._3)
)

combineByKey : la forme générale

Les trois précédentes en sont des cas particuliers. Elle demande trois fonctions :

def combineByKey[C](
  creerCombineur: V => C,               // première valeur rencontrée pour une clé
  fusionnerValeur: (C, V) => C,         // valeur suivante dans la même partition
  fusionnerCombineurs: (C, C) => C      // entre partitions
): RDD[(K, C)]

La même moyenne :

val moyennes = notes
  .combineByKey(
    (note: Int) => (note, 1),                              // créer
    (acc: (Int, Int), note: Int) => (acc._1 + note, acc._2 + 1),  // fusionner une valeur
    (a: (Int, Int), b: (Int, Int)) => (a._1 + b._1, a._2 + b._2)  // fusionner deux accs
  )
  .mapValues { case (somme, compte) => somme.toDouble / compte }

La différence avec aggregateByKey est subtile mais utile : au lieu d’une valeur initiale fixe, on fournit une fonction qui construit l’accumulateur à partir de la première valeur rencontrée. C’est indispensable quand l’accumulateur ne peut pas avoir de valeur neutre naturelle — par exemple si vous accumulez dans une structure qui doit être initialisée avec une vraie donnée.

Dans 95 % des cas, aggregateByKey suffit et se lit mieux. combineByKey reste la porte de sortie quand elle ne suffit pas.

Et en DataFrame ? Beaucoup plus court

Puisque ces méthodes appartiennent à l’API RDD, la question se pose : en DataFrame, une moyenne par clé s’écrit ainsi :

import org.apache.spark.sql.functions._

notes.toDF("prenom", "note")
  .groupBy("prenom")
  .agg(avg("note").as("moyenne"), count("*").as("nb"))

Une ligne, et Catalyst optimise tout — avec le même combine local, sans que vous ayez à l’exprimer. C’est l’illustration directe de ce qui est développé dans RDD ou DataFrame : la famille aggregateByKey reste indispensable pour comprendre le mécanisme et pour les cas sans équivalent SQL, mais en production, l’agrégation DataFrame gagne presque toujours.

Ce qu’il faut retenir

reduceByKey bloque dès que le résultat change de type — c’est la limite, et c’est pour la franchir que la famille existe. foldByKey ajoute un élément neutre. aggregateByKey autorise un accumulateur d’un type différent, et couvre à peu près tout : moyennes, min/max simultanés, ensembles distincts, statistiques. combineByKey est la forme générale, à garder pour les accumulateurs sans valeur neutre.

Le critère qui compte dans tous les cas : un accumulateur de taille bornée. C’est ce qui garantit que le shuffle transporte quelques octets par clé au lieu de toutes les valeurs — la différence exacte entre ces méthodes et un groupByKey qui finit par tomber en mémoire.