Lancer des lectures de données JDBC simultanées

Avez-vous déjà suivi le processus d'implémentation de Spark dans votre projet, en déterminant le nombre optimal de partitions aléatoires, l'allocation de mémoire pour les instances de pilote et d'exécuteur, le nombre de cœurs d'exécuteur et toutes ces choses amusantes juste pour lire des données à partir d'une source JDBC comme celle-ci ?

jdbcDF = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("query", query) \
    .load()

Quel est le problème avec le code ci-dessus ? Eh bien… c'est très lent :) Si vous exécutez une requête SQL dans Spark comme dans l'exemple ci-dessus, vous n'utiliserez qu'un seul thread (vous pouvez voir dans l'interface utilisateur Spark qu'une seule tâche s'exécute lors de l'exécution du code) et le dataframe créé sera avoir une partition.

La plupart d'entre nous utilisent Spark dans les pipelines ETL/ELT pour le traitement parallèle, donc lire des données à partir d'une source JDBC en utilisant un seul thread n'est probablement pas ce que nous voulons. Sans oublier que nous devrons repartitionner la dataframe après l'exécution de la requête car les données sont désormais stockées dans 1 partition. Alors, comment lisons-nous réellement les données simultanément dans Spark ?

# pass the original query to dbtable as a subquery 
final_query = f'({query}) as q'

# depends on your usecase, how many executor cores you have in the cluster
# what do you plan to do with the dataframe downstream, data size etc
partitionsCount = 100

# I'm hardcoding the date values here, but I'll show below how to identify
# these values at run time if you don't know them upfront
min_date = '2020-01-01'
max_date = '2022-12-31'

jdbcDF = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("dbtable", final_query) \
    .option("numPartitions", partitionsCount) \
    .option("partitionColumn", "date") \
    .option("lowerBound", f"{min_date}") \
    .option("upperBound", f"{max_date}") \
    .load()

1. dbtable :
   D'après la documentation officielle de Spark, si nous voulons lire des données simultanément et les partitionner en fonction de nos besoins, nous ne pouvons plus utiliser l' option de requête et devons plutôt passer notre requête SQL via l' option dbtable . La seule différence entre ces deux options (autre que le nom) est que la requête doit maintenant être entre parenthèses et avoir un alias — en gros, passez la requête d'origine en tant que sous-requête.
2. numPartitions :
   nombre maximal de partitions pouvant être utilisées pour le parallélisme dans la lecture et l'écriture de tables. Cela détermine également le nombre maximal de connexions JDBC simultanées. Si le nombre de partitions à écrire dépasse cette limite, on le diminue jusqu'à cette limite en appelant coalesce(numPartitions) avant d'écrire. ( source ) J'ai le mien défini sur 100 dans le cluster EMR, mais une bonne pratique serait de définir le numéro de partition entre 1 et 4 fois le nombre de cœurs. Ce n'est pas un nombre gravé dans la pierre, alors testez toujours avec des valeurs différentes pour voir ce qui fonctionne pour vous.
3. partitionColumn :
   C'est probablement le paramètre le plus important et celui qui fera la plus grande différence dans le temps d'exécution de vos requêtes. Peu de choses à garder à l'esprit :
   - la colonne en question doit être numérique, date ou horodatage
   - pour les meilleures performances possibles, les valeurs de la colonne doivent être réparties aussi uniformément que possible et avoir une cardinalité élevée. Nous voulons éviter les colonnes de données asymétriques - je développerai cela ci-dessous.
   - si vous avez plusieurs colonnes avec les propriétés ci-dessus, sélectionnez la colonne indexée.
4. lowerbound, upperbound :
   Notez que lowerBound et upperBound sont juste utilisés pour décider de la foulée de la partition, pas pour filtrer les lignes de la table. Ainsi, toutes les lignes de la table seront partitionnées et renvoyées. ( source ) Ci-dessous, vous verrez comment Spark génère chaque partition. C'est pourquoi vous ne voulez pas d'une colonne à faible cardinalité comme colonne de partition, mais je vais donner un exemple ci-dessous au cas où cela n'aurait pas de sens maintenant.

Imaginez que la colonne que vous souhaitez utiliser pour le partitionnement n'ait que des valeurs 0 et 1 - un exemple extrême en effet, mais un scénario courant en Data Engineering où les valeurs booléennes True et False sont converties en valeurs numériques.

Comme vous n'avez que des 0 et des 1, vous utilisez 0 comme limite inférieure et 1 comme limite supérieure. En fonction de la façon dont les partitions sont générées - voir l'image ci-dessus - toutes les lignes avec la valeur 0 seront poussées vers la première partition (la clause WHERE de la première partition est la seule clause où 0 n'est pas filtré) et toutes les lignes avec la valeur 1 sera poussé vers la dernière partition (la clause WHERE de la dernière partition est la seule clause où 1 n'est pas filtré) .

Ainsi, quel que soit le nombre de partitions que vous créez, seules deux partitions seront peuplées et seuls deux cœurs d'exécuteur feront tout le travail, les autres seront en état d'inactivité s'il n'y a pas d'autre tâche prévue pour eux. À cause de cela, les performances de votre requête en prendront un coup.

Asymétrie des données

Imaginons que nous ayons un ensemble de données avec 20 millions de lignes et 30 partitions, les limites inférieure et supérieure étant 2020–01–01 et 2022–12–31. Pour les besoins de cet article, j'utiliserai une session Spark avec 30 cœurs d'exécuteur (ce qui signifie que nous pouvons exécuter jusqu'à 30 tâches en même temps). Gardez à l'esprit que Spark attribue une tâche par partition, de sorte que chaque partition sera traitée par un cœur d'exécuteur.

Chaque année, la valeur des données sera divisée en 10 partitions et étant donné qu'une tâche sera affectée à une partition, chaque tâche affectée aux années 2020 et 2021 sera responsable du traitement de 100 000 lignes, mais chaque tâche affectée à l'année 2022 devra traiter 1,8 M de lignes. Cela représente une augmentation de 1 700 % des données qui doivent être traitées par une seule tâche.

Le résultat est que les 20 premières tâches finiront de traiter leur partition assignée en un rien de temps et après cela, les 20 cœurs d'exécuteur affectés aux tâches terminées resteront inactifs (s'il n'y a pas d'autre travail où les cœurs d'exécuteur inactifs peuvent être utilisés) jusqu'au dernier 10 tâches terminent le traitement de 1 700 % d'enregistrements supplémentaires. Il s'agit d'un gaspillage de ressources et d'une augmentation des coûts d'exploitation, en particulier si vous utilisez Spark dans un cluster EMR, Glue Job ou Databricks.

Comment résoudre ce problème ? La solution la plus simple : utilisez une colonne différente pour le partitionnement. Mais s'il est essentiel d'utiliser cette colonne, nous pouvons utiliser deux requêtes et deux dataframes au lieu d'un. Vous aurez toujours des problèmes avec des données asymétriques en aval (un problème récurrent dans les systèmes de calcul parallèles), mais nous nous attaquons à un problème à la fois.

La première requête chargera les deux premières années (les 30 cœurs exécuteurs seront répartis entre 2020 et 2021), puis la deuxième requête chargera la dernière année (donc tous les 30 cœurs exécuteurs seront désormais utilisés pour 2022). Pour cette raison, une tâche donnée associée à 2022 sera désormais responsable du traitement de 600 000 lignes au lieu de 1,8 million de lignes. Cela représente une diminution de 66 % des lignes qui doivent être traitées par un seul cœur d'exécuteur par rapport à l'implémentation d'origine. En plus de cela, nous utilisons maintenant tous les cœurs d'exécuteur de notre session Spark au lieu de les laisser inactifs.

Que se passe-t-il si je n'ai pas de colonne numérique, de date ou d'horodatage, ou si la colonne contient des données faussées ?

Pas de problème. En tant qu'ingénieurs de données, l'utilisation d'un CTE et d'une fonction de fenêtre pour générer une colonne numérique row_number()devrait être un jeu d'enfant. Une surcharge sera introduite lors de la génération de la nouvelle colonne, mais même avec cette surcharge, la requête finale sera exécutée beaucoup plus rapidement en utilisant la colonne nouvellement créée plutôt que d'essayer de charger toutes les données dans une partition à l'aide d'un cœur d'exécuteur.

with cte as
(
    SELECT f1, f2, f3
    FROM table
)
select *, row_number() over (order by f1) as rn from cte

original_query = "SELECT f1, f2, f3 FROM table"

final_query = f'(with cte as ({original_query}) select *, row_number() over (order by f1) as rn from cte) as q'

Nous exécutons deux requêtes au lieu d'une :
- la première requête sera exécutée pour trouver les valeurs min et max de la colonne que nous utiliserons comme partitionColumn -
dans la deuxième requête, nous passerons les résultats de la première requête en arguments lowerbound et upperbound .

query_min_max = f"SELECT min(date), max(date) from ({query}) q"
final_query = f"({query}) as q"

df_min_max = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("query", query_min_max) \
    .load()

min = df_min_max.first()["min"]
max = df_min_max.first()["max"]

jdbcDF = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("dbtable", final_query) \
    .option("numPartitions", partitionsCount) \
    .option("partitionColumn", "date") \
    .option("lowerBound", f"{min}") \
    .option("upperBound", f"{max}") \
    .load()

• Si vous suivez le processus d'implémentation de Spark dans votre application, essayez de l'utiliser au maximum et lisez les données simultanément. L'exécution d'une requête SQL sur spark.read() comme indiqué dans la documentation JDBC de Spark placera toutes les données dans une partition et n'utilisera qu'un seul cœur d'exécuteur, quel que soit le nombre de cœurs que vous avez configurés dans votre session Spark.
• Si vous utilisez Spark aujourd'hui pour lire des données à partir de sources de données JDBC en utilisant l'approche à thread unique, essayez de refactoriser vos requêtes et d'améliorer leurs performances, vous pourriez même obtenir une promotion pendant que vous y êtes :)
• Au bout du compte, le temps c'est de l'argent. Si vous vous inquiétez de l'effort de réécriture de vos requêtes, je peux parler d'expérience - certaines des requêtes Spark que j'ai refactorisées à l'aide des techniques décrites ci-dessus s'exécutent désormais > 90 % plus rapidement qu'auparavant. Vous pouvez imaginer quel accueil positif cette amélioration a eu dans mon équipe.

gabriel