UDF Spark : pourquoi la vôtre est dix fois plus lente

Vous écrivez une petite fonction Scala, vous l’enregistrez comme UDF, elle fait exactement ce qu’il faut. Puis le job passe de quatre minutes à trente-cinq. Ce n’est pas votre logique qui est lente — c’est que vous venez de rendre aveugle l’optimiseur de Spark sur cette colonne.

Une UDF n’est pas interdite. Mais elle a un coût précis, mesurable, et évitable dans la grande majorité des cas. Voici lequel, et comment s’en passer.

Ce que Catalyst voit — et ne voit plus

Avec une fonction native, Spark lit l’expression :

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

df.withColumn("majuscule", upper($"nom"))

Catalyst sait que c’est upper, connaît son comportement sur les null, peut la fusionner avec les opérations voisines dans une seule boucle compilée, et peut décider de pousser le filtre associé jusqu’au fichier Parquet.

Avec une UDF :

val majusculeUdf = udf((s: String) => s.toUpperCase)

df.withColumn("majuscule", majusculeUdf($"nom"))

Catalyst voit une boîte noire. Il ne peut plus rien en déduire. Concrètement, vous perdez quatre choses :

Le troisième point est souvent le plus coûteux : un filtre qui aurait éliminé 99 % des données au niveau du fichier doit maintenant charger 100 % des lignes pour les évaluer une par une.

Sur Databricks, la facture est double : une UDF fait aussi retomber l’étape sur Spark JVM au lieu du moteur Photon, que vous payez pourtant. Le détail est dans Snowflake vs Databricks : Spark, Photon et le moteur.

Les remplacements que vous cherchez probablement

La plupart des UDF écrites en production ont un équivalent natif. Voici les cas les plus fréquents :

Exemple typique, avant :

val categoriser = udf((montant: Double) =>
  if (montant > 1000) "grand" else if (montant > 100) "moyen" else "petit"
)
df.withColumn("categorie", categoriser($"montant"))

Après — même résultat, entièrement optimisable :

df.withColumn("categorie",
  when($"montant" > 1000, "grand")
    .when($"montant" > 100, "moyen")
    .otherwise("petit")
)

Depuis Spark 3, les fonctions d’ordre supérieur couvrent aussi les cas qui obligeaient historiquement à écrire une UDF sur les tableaux :

// Doubler chaque élément d'un tableau, sans UDF
df.withColumn("doubles", transform($"valeurs", x => x * 2))

// Filtrer un tableau
df.withColumn("positifs", filter($"valeurs", x => x > 0))

Quand l’UDF est vraiment inévitable

Il reste des cas légitimes, et il faut savoir les reconnaître :

• Une bibliothèque métier que vous devez appeler (géocodage, modèle ML embarqué, parseur d’un format propriétaire).
• Un algorithme réellement non exprimable en fonctions natives — rare, mais ça existe.
• Du code déjà validé et audité qu’il serait risqué de réécrire pour gagner quelques minutes.

Dans ce cas, réduisez les dégâts :

// 1. Typer explicitement le retour : évite une inférence coûteuse
val monUdf = udf[String, String]((s: String) => traiter(s))

// 2. Filtrer AVANT l'UDF, jamais après
df.filter($"actif" === true)      // ← pushdown préservé
  .withColumn("resultat", monUdf($"donnee"))

// 3. Gérer les null soi-même
val sûr = udf((s: String) => if (s == null) null else traiter(s))

L’ordre du point 2 est le levier le plus important : appliquer l’UDF sur les 2 % de lignes qui survivent au filtre plutôt que sur 100 % change tout, et ne demande qu’un réarrangement de lignes.

from pyspark.sql.functions import pandas_udf
import pandas as pd

@pandas_udf("double")
def normaliser(s: pd.Series) -> pd.Series:
    return (s - s.mean()) / s.std()

df.withColumn("norm", normaliser("valeur"))

Mesurer, plutôt que supposer

Le plan explain montre littéralement l’endroit où l’optimisation s’arrête :

df.withColumn("x", monUdf($"col")).explain(true)
// == Physical Plan ==
// *(1) Project [..., UDF(col#12) AS x#20]
//                    ^^^ boîte noire : pas de codegen sur cette expression

Comparez avec la version native, où l’expression apparaît en clair (upper(col#12), CASE WHEN ...) et est intégrée dans l’étage compilé. Et si vous hésitez entre deux écritures, mesurez sur un échantillon représentatif — pas sur dix lignes, où tout paraît instantané.

Ce qu’il faut retenir

Une UDF échange de la lisibilité contre de l’optimisation, et le taux de change est mauvais : vous perdez le codegen, le format binaire, le pushdown et la gestion des null, souvent pour du code qui existait déjà dans org.apache.spark.sql.functions. Le réflexe à installer : chercher la fonction native d’abord, l’UDF seulement quand on peut expliquer pourquoi aucune ne convient.

Et quand l’UDF est inévitable, deux gestes suffisent à limiter la casse : filtrer en amont pour réduire le nombre d’appels, et gérer les null explicitement. Le contexte général de cette optimisation — pourquoi Catalyst est le vrai moteur de performance de Spark — est développé dans RDD ou DataFrame.