Le moyen le plus rapide de charger des données dans PostgreSQL à l'aide de Python

Référence

Au moment de la rédaction, l'API des bières ne contient que 325 bières. Pour travailler sur un grand ensemble de données, nous le dupliquons 100 fois et le stockons en mémoire. L'ensemble de données résultant contient 32 500 bières :

>>> beers = list(iter_beers_from_api()) * 100
>>> len(beers)
32,500

Pour imiter une API distante, nos fonctions accepteront des itérateurs similaires à la valeur de retour de iter_beers_from_api :

def process(beers: Iterator[Dict[str, Any]])) -> None:
 # Process beers...

Pour le benchmark, nous allons importer les données de la bière dans la base de données. Pour éliminer les influences externes telles que le réseau, nous récupérons les données de l'API à l'avance et les diffusons localement.

Pour obtenir un timing précis, nous "truquons" l'API distante :

>>> beers = list(iter_beers_from_api()) * 100
>>> process(beers)

Dans une situation réelle, vous utiliseriez la fonction iter_beers_from_api directement :

>>> process(iter_beers_from_api())

Nous sommes maintenant prêts à commencer !

Insérer des lignes une par une

Pour établir une ligne de base, nous commençons par l'approche la plus simple, insérez les lignes une par une :

@profile
def insert_one_by_one(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 for beer in beers:
 cursor.execute("""
 INSERT INTO staging_beers VALUES (
 %(id)s,
 %(name)s,
 %(tagline)s,
 %(first_brewed)s,
 %(description)s,
 %(image_url)s,
 %(abv)s,
 %(ibu)s,
 %(target_fg)s,
 %(target_og)s,
 %(ebc)s,
 %(srm)s,
 %(ph)s,
 %(attenuation_level)s,
 %(brewers_tips)s,
 %(contributed_by)s,
 %(volume)s
 );
 """, {
 **beer,
 'first_brewed': parse_first_brewed(beer['first_brewed']),
 'volume': beer['volume']['value'],
 })

Notez que lorsque nous itérons les bières, nous transformons le first_brewed à un datetime.date et extrait la valeur de volume du volume imbriqué champ.

L'exécution de cette fonction produit la sortie suivante :

>>> insert_one_by_one(connection, beers)
insert_one_by_one()
Time 128.8
Memory 0.08203125

La fonction a pris 129 secondes pour importer 32 000 lignes. Le profileur de mémoire montre que la fonction a consommé très peu de mémoire.

Intuitivement, insérer des lignes une par une ne semble pas très efficace. Le changement constant de contexte entre le programme et la base de données doit le ralentir.

Exécuter plusieurs

Psycopg2 fournit un moyen d'insérer plusieurs lignes à la fois en utilisant executemany . À partir de la documentation :

Cela semble prometteur !

Essayons d'importer les données en utilisant executemany :

@profile
def insert_executemany(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)

 all_beers = [{
 **beer,
 'first_brewed': parse_first_brewed(beer['first_brewed']),
 'volume': beer['volume']['value'],
 } for beer in beers]

 cursor.executemany("""
 INSERT INTO staging_beers VALUES (
 %(id)s,
 %(name)s,
 %(tagline)s,
 %(first_brewed)s,
 %(description)s,
 %(image_url)s,
 %(abv)s,
 %(ibu)s,
 %(target_fg)s,
 %(target_og)s,
 %(ebc)s,
 %(srm)s,
 %(ph)s,
 %(attenuation_level)s,
 %(brewers_tips)s,
 %(contributed_by)s,
 %(volume)s
 );
 """, all_beers)

La fonction ressemble beaucoup à la fonction précédente et les transformations sont les mêmes. La principale différence ici est que nous transformons d'abord toutes les données en mémoire, puis que nous les importons dans la base de données.

L'exécution de cette fonction produit la sortie suivante :

>>> insert_executemany(connection, beers)
insert_executemany()
Time 124.7
Memory 2.765625

C'est décevant. Le timing est juste un peu meilleur, mais la fonction consomme maintenant 2,7 Mo de mémoire.

Pour mettre l'utilisation de la mémoire en perspective, un fichier JSON contenant uniquement les données que nous importons pèse 25 Mo sur disque. Compte tenu de la proportion, l'utilisation de cette méthode pour importer un fichier de 1 Go nécessitera 110 Mo de mémoire.

Exécuter plusieurs depuis l'itérateur

La méthode précédente consommait beaucoup de mémoire car les données transformées étaient stockées en mémoire avant d'être traitées par psycopg.

Voyons si nous pouvons utiliser un itérateur pour éviter de stocker les données en mémoire :

@profile
def insert_executemany_iterator(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 cursor.executemany("""
 INSERT INTO staging_beers VALUES (
 %(id)s,
 %(name)s,
 %(tagline)s,
 %(first_brewed)s,
 %(description)s,
 %(image_url)s,
 %(abv)s,
 %(ibu)s,
 %(target_fg)s,
 %(target_og)s,
 %(ebc)s,
 %(srm)s,
 %(ph)s,
 %(attenuation_level)s,
 %(brewers_tips)s,
 %(contributed_by)s,
 %(volume)s
 );
 """, ({
 **beer,
 'first_brewed': parse_first_brewed(beer['first_brewed']),
 'volume': beer['volume']['value'],
 } for beer in beers))

La différence ici est que les données transformées sont "diffusées" en executemany en utilisant un itérateur.

Cette fonction produit le résultat suivant :

>>> insert_executemany_iterator(connection, beers)
insert_executemany_iterator()
Time 129.3
Memory 0.0

Notre solution de "streaming" a fonctionné comme prévu et nous avons réussi à ramener la mémoire à zéro. Cependant, le timing reste à peu près le même, même par rapport à la méthode un par un.

Exécuter le lot

La documentation de psycopg contient une note très intéressante sur executemany dans la section "aides à l'exécution rapide" :

Donc, nous nous sommes trompés depuis le début !

La fonction juste en dessous de cette section est execute_batch :

Implémentons la fonction de chargement en utilisant execute_batch :

import psycopg2.extras

@profile
def insert_execute_batch(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)

 all_beers = [{
 **beer,
 'first_brewed': parse_first_brewed(beer['first_brewed']),
 'volume': beer['volume']['value'],
 } for beer in beers]

 psycopg2.extras.execute_batch(cursor, """
 INSERT INTO staging_beers VALUES (
 %(id)s,
 %(name)s,
 %(tagline)s,
 %(first_brewed)s,
 %(description)s,
 %(image_url)s,
 %(abv)s,
 %(ibu)s,
 %(target_fg)s,
 %(target_og)s,
 %(ebc)s,
 %(srm)s,
 %(ph)s,
 %(attenuation_level)s,
 %(brewers_tips)s,
 %(contributed_by)s,
 %(volume)s
 );
 """, all_beers)

Exécution de la fonction :

>>> insert_execute_batch(connection, beers)
insert_execute_batch()
Time 3.917
Memory 2.50390625

Ouah! C'est un énorme bond en avant. La fonction terminée en un peu moins de 4 secondes. C'est environ 33 fois plus rapide que les 129 secondes avec lesquelles nous avons commencé.

Exécuter le lot à partir de l'itérateur

La fonction execute_batch utilisé moins de mémoire que executemany fait pour les mêmes données. Essayons d'éliminer la mémoire en "streamant" les données dans execute_batch en utilisant un itérateur :

@profile
def insert_execute_batch_iterator(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)

 iter_beers = ({
 **beer,
 'first_brewed': parse_first_brewed(beer['first_brewed']),
 'volume': beer['volume']['value'],
 } for beer in beers)

 psycopg2.extras.execute_batch(cursor, """
 INSERT INTO staging_beers VALUES (
 %(id)s,
 %(name)s,
 %(tagline)s,
 %(first_brewed)s,
 %(description)s,
 %(image_url)s,
 %(abv)s,
 %(ibu)s,
 %(target_fg)s,
 %(target_og)s,
 %(ebc)s,
 %(srm)s,
 %(ph)s,
 %(attenuation_level)s,
 %(brewers_tips)s,
 %(contributed_by)s,
 %(volume)s
 );
 """, iter_beers)

Exécution de la fonction

>>> insert_execute_batch_iterator(connection, beers)
insert_execute_batch_iterator()
Time 4.333
Memory 0.2265625

Nous avons à peu près le même temps, mais avec moins de mémoire.

Exécuter un lot à partir de l'itérateur avec la taille de la page

Lors de la lecture de la documentation pour execute_batch , l'argument page_size a attiré mon attention :

La documentation indiquait précédemment que la fonction fonctionnait mieux car elle effectuait moins d'allers-retours vers la base de données. Si tel est le cas, une taille de page plus grande devrait réduire le nombre d'allers-retours et entraîner un temps de chargement plus rapide.

Ajoutons un argument pour la taille de la page à notre fonction afin que nous puissions expérimenter :

@profile
def insert_execute_batch_iterator(
 connection,
 beers: Iterator[Dict[str, Any]],
 page_size: int = 100,
) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)

 iter_beers = ({
 **beer,
 'first_brewed': parse_first_brewed(beer['first_brewed']),
 'volume': beer['volume']['value'],
 } for beer in beers)

 psycopg2.extras.execute_batch(cursor, """
 INSERT INTO staging_beers VALUES (
 %(id)s,
 %(name)s,
 %(tagline)s,
 %(first_brewed)s,
 %(description)s,
 %(image_url)s,
 %(abv)s,
 %(ibu)s,
 %(target_fg)s,
 %(target_og)s,
 %(ebc)s,
 %(srm)s,
 %(ph)s,
 %(attenuation_level)s,
 %(brewers_tips)s,
 %(contributed_by)s,
 %(volume)s
 );
 """, iter_beers, page_size=page_size)

La taille de page par défaut est de 100. Comparons différentes valeurs et comparons les résultats :

>>> insert_execute_batch_iterator(connection, iter(beers), page_size=1)
insert_execute_batch_iterator(page_size=1)
Time 130.2
Memory 0.0

>>> insert_execute_batch_iterator(connection, iter(beers), page_size=100)
insert_execute_batch_iterator(page_size=100)
Time 4.333
Memory 0.0

>>> insert_execute_batch_iterator(connection, iter(beers), page_size=1000)
insert_execute_batch_iterator(page_size=1000)
Time 2.537
Memory 0.2265625

>>> insert_execute_batch_iterator(connection, iter(beers), page_size=10000)
insert_execute_batch_iterator(page_size=10000)
Time 2.585
Memory 25.4453125

Nous avons obtenu des résultats intéressants, décomposons-les :

• 1 :Les résultats sont similaires aux résultats que nous avons obtenus en insérant les lignes une par une.
• 100 :il s'agit du page_size par défaut , les résultats sont donc similaires à ceux de notre benchmark précédent.
• 1 000 :la synchronisation est ici environ 40 % plus rapide et la mémoire est faible.
• 10 000 :la synchronisation n'est pas beaucoup plus rapide qu'avec une taille de page de 1 000, mais la mémoire est nettement plus élevée.

Les résultats montrent qu'il existe un compromis entre la mémoire et la vitesse. Dans ce cas, il semble que le sweet spot corresponde à une taille de page de 1 000.

Exécuter les valeurs

Les joyaux de la documentation de psycopg ne se terminent pas par execute_batch . En parcourant la documentation, une autre fonction appelée execute_values a attiré mon attention :

La fonction execute_values fonctionne en générant une énorme liste de VALEURS à la requête.

Essayons :

import psycopg2.extras

@profile
def insert_execute_values(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 psycopg2.extras.execute_values(cursor, """
 INSERT INTO staging_beers VALUES %s;
 """, [(
 beer['id'],
 beer['name'],
 beer['tagline'],
 parse_first_brewed(beer['first_brewed']),
 beer['description'],
 beer['image_url'],
 beer['abv'],
 beer['ibu'],
 beer['target_fg'],
 beer['target_og'],
 beer['ebc'],
 beer['srm'],
 beer['ph'],
 beer['attenuation_level'],
 beer['brewers_tips'],
 beer['contributed_by'],
 beer['volume']['value'],
 ) for beer in beers])

Importer des bières avec la fonction :

>>> insert_execute_values(connection, beers)
insert_execute_values()
Time 3.666
Memory 4.50390625

Donc, dès la sortie de la boîte, nous obtenons une légère accélération par rapport à execute_batch . Cependant, la mémoire est légèrement supérieure.

Exécuter les valeurs de l'itérateur

Tout comme nous l'avons fait auparavant, pour réduire la consommation de mémoire, nous essayons d'éviter de stocker des données en mémoire en utilisant un itérateur au lieu d'une liste :

@profile
def insert_execute_values_iterator(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 psycopg2.extras.execute_values(cursor, """
 INSERT INTO staging_beers VALUES %s;
 """, ((
 beer['id'],
 beer['name'],
 beer['tagline'],
 parse_first_brewed(beer['first_brewed']),
 beer['description'],
 beer['image_url'],
 beer['abv'],
 beer['ibu'],
 beer['target_fg'],
 beer['target_og'],
 beer['ebc'],
 beer['srm'],
 beer['ph'],
 beer['attenuation_level'],
 beer['brewers_tips'],
 beer['contributed_by'],
 beer['volume']['value'],
 ) for beer in beers))

L'exécution de la fonction a produit les résultats suivants :

>>> insert_execute_values_iterator(connection, beers)
insert_execute_values_iterator()
Time 3.677
Memory 0.0

Ainsi, le timing est presque le même, mais la mémoire est de retour à zéro.

Exécuter les valeurs de l'itérateur avec la taille de la page

Tout comme execute_batch , la fonction execute_values accepter également un page_size argument :

@profile
def insert_execute_values_iterator(
 connection,
 beers: Iterator[Dict[str, Any]],
 page_size: int = 100,
) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 psycopg2.extras.execute_values(cursor, """
 INSERT INTO staging_beers VALUES %s;
 """, ((
 beer['id'],
 beer['name'],
 beer['tagline'],
 parse_first_brewed(beer['first_brewed']),
 beer['description'],
 beer['image_url'],
 beer['abv'],
 beer['ibu'],
 beer['target_fg'],
 beer['target_og'],
 beer['ebc'],
 beer['srm'],
 beer['ph'],
 beer['attenuation_level'],
 beer['brewers_tips'],
 beer['contributed_by'],
 beer['volume']['value'],
 ) for beer in beers), page_size=page_size)

Exécution avec différentes tailles de page :

>>> insert_execute_values_iterator(connection, iter(beers), page_size=1)
insert_execute_values_iterator(page_size=1)
Time 127.4
Memory 0.0

>>> insert_execute_values_iterator(connection, iter(beers), page_size=100)
insert_execute_values_iterator(page_size=100)
Time 3.677
Memory 0.0

>>> insert_execute_values_iterator(connection, iter(beers), page_size=1000)
insert_execute_values_iterator(page_size=1000)
Time 1.468
Memory 0.0

>>> insert_execute_values_iterator(connection, iter(beers), page_size=10000)
insert_execute_values_iterator(page_size=10000)
Time 1.503
Memory 2.25

Tout comme execute_batch , nous voyons un compromis entre la mémoire et la vitesse. Ici aussi, le sweet spot se situe autour de la taille de page 1000. Cependant, en utilisant execute_values nous avons obtenu des résultats environ 20 % plus rapides par rapport à la même taille de page en utilisant execute_batch .

Copier

La documentation officielle de PostgreSQL contient une section entière sur le remplissage d'une base de données. Selon la documentation, la meilleure façon de charger des données dans une base de données est d'utiliser le copy commande.

Pour utiliser copy à partir de Python, psycopg fournit une fonction spéciale appelée copy_from . Le copy La commande nécessite un fichier CSV. Voyons si nous pouvons transformer nos données en CSV et les charger dans la base de données en utilisant copy_from :

import io

def clean_csv_value(value: Optional[Any]) -> str:
 if value is None:
 return r'\N'
 return str(value).replace('\n', '\\n')

@profile
def copy_stringio(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 csv_file_like_object = io.StringIO()
 for beer in beers:
 csv_file_like_object.write('|'.join(map(clean_csv_value, (
 beer['id'],
 beer['name'],
 beer['tagline'],
 parse_first_brewed(beer['first_brewed']),
 beer['description'],
 beer['image_url'],
 beer['abv'],
 beer['ibu'],
 beer['target_fg'],
 beer['target_og'],
 beer['ebc'],
 beer['srm'],
 beer['ph'],
 beer['attenuation_level'],
 beer['contributed_by'],
 beer['brewers_tips'],
 beer['volume']['value'],
 ))) + '\n')
 csv_file_like_object.seek(0)
 cursor.copy_from(csv_file_like_object, 'staging_beers', sep='|')

Décomposons-le :

• clean_csv_value :Transforme une seule valeur
  • Échappez aux nouvelles lignes  :certains des champs de texte incluent des retours à la ligne, nous échappons donc à \n -> \\n .
  • Les valeurs vides sont transformées en \N :La chaîne "\N" est la chaîne par défaut utilisée par PostgreSQL pour indiquer NULL dans COPY (cela peut être changé en utilisant le NULL option).
• csv_file_like_object :Générer un fichier comme un objet en utilisant io.StringIO . Un StringIO objet contient une chaîne qui peut être utilisée comme un fichier. Dans notre cas, un fichier CSV.
• csv_file_like_object.write :Transformer une bière en ligne CSV
  • Transformez les données :transformations sur first_brewed et volume sont exécutés ici.
  • Choisissez un délimiteur :Certains des champs du jeu de données contiennent du texte libre avec des virgules. Pour éviter les conflits, nous choisissons "|" comme délimiteur (une autre option consiste à utiliser QUOTE ).

Voyons maintenant si tout ce travail acharné a porté ses fruits :

>>> copy_stringio(connection, beers)
copy_stringio()
Time 0.6274
Memory 99.109375

Le copy commande est la plus rapide que nous ayons vue jusqu'à présent ! Utilisation de COPY , le processus terminé en moins d'une seconde. Cependant, il semble que cette méthode soit beaucoup plus coûteuse en termes d'utilisation de la mémoire. La fonction consomme 99 Mo, soit plus du double de la taille de notre fichier JSON sur le disque.

Copier les données d'un itérateur de chaîne

L'un des principaux inconvénients de l'utilisation de la copie avec StringIO est que le fichier entier est créé en mémoire. Et si au lieu de créer le fichier entier en mémoire, nous créons un objet semblable à un fichier qui agira comme un tampon entre la source distante et le COPY commande. Le tampon consommera JSON via l'itérateur, nettoiera et transformera les données, et produira un CSV propre.

Inspirés par cette réponse de débordement de pile, nous avons créé un objet qui se nourrit d'un itérateur et fournit une interface de type fichier :

from typing import Iterator, Optional
import io

class StringIteratorIO(io.TextIOBase):
 def __init__(self, iter: Iterator[str]):
 self._iter = iter
 self._buff = ''

 def readable(self) -> bool:
 return True

 def _read1(self, n: Optional[int] = None) -> str:
 while not self._buff:
 try:
 self._buff = next(self._iter)
 except StopIteration:
 break
 ret = self._buff[:n]
 self._buff = self._buff[len(ret):]
 return ret

 def read(self, n: Optional[int] = None) -> str:
 line = []
 if n is None or n < 0:
 while True:
 m = self._read1()
 if not m:
 break
 line.append(m)
 else:
 while n > 0:
 m = self._read1(n)
 if not m:
 break
 n -= len(m)
 line.append(m)
 return ''.join(line)

Pour montrer comment cela fonctionne, voici comment un objet semblable à un fichier CSV peut être généré à partir d'une liste de nombres :

>>> gen = (f'{i},{i**2}\n' for i in range(3))
>>> gen
<generator object <genexpr> at 0x7f58bde7f5e8>
>>> f = StringIteratorIO(gen)
>>> print(f.read())
0,0
1,1
2,4

Notez que nous avons utilisé f comme un fichier. En interne, il a récupéré les lignes de gen uniquement lorsque son tampon de ligne interne était vide.

La fonction de chargement utilisant StringIteratorIO ressemble à ceci :

@profile
def copy_string_iterator(connection, beers: Iterator[Dict[str, Any]]) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 beers_string_iterator = StringIteratorIO((
 '|'.join(map(clean_csv_value, (
 beer['id'],
 beer['name'],
 beer['tagline'],
 parse_first_brewed(beer['first_brewed']).isoformat(),
 beer['description'],
 beer['image_url'],
 beer['abv'],
 beer['ibu'],
 beer['target_fg'],
 beer['target_og'],
 beer['ebc'],
 beer['srm'],
 beer['ph'],
 beer['attenuation_level'],
 beer['brewers_tips'],
 beer['contributed_by'],
 beer['volume']['value'],
 ))) + '\n'
 for beer in beers
 ))
 cursor.copy_from(beers_string_iterator, 'staging_beers', sep='|')

La principale différence est que le fichier CSV beers est consommé à la demande et que les données ne sont pas stockées en mémoire après leur utilisation.

Exécutons la fonction et voyons les résultats :

>>> copy_string_iterator(connection, beers)
copy_string_iterator()
Time 0.4596
Memory 0.0

Super! Le temps est faible et la mémoire est de retour à zéro.

Copier les données d'un itérateur de chaîne avec une taille de tampon

Dans une tentative de presser une dernière goutte de performance, nous remarquons que tout comme page_size , le copy La commande accepte également un argument similaire appelé size :

Ajoutons un size argument de la fonction :

@profile
def copy_string_iterator(connection, beers: Iterator[Dict[str, Any]], size: int = 8192) -> None:
 with connection.cursor() as cursor:
 create_staging_table(cursor)
 beers_string_iterator = StringIteratorIO((
 '|'.join(map(clean_csv_value, (
 beer['id'],
 beer['name'],
 beer['tagline'],
 parse_first_brewed(beer['first_brewed']).isoformat(),
 beer['description'],
 beer['image_url'],
 beer['abv'],
 beer['ibu'],
 beer['target_fg'],
 beer['target_og'],
 beer['ebc'],
 beer['srm'],
 beer['ph'],
 beer['attenuation_level'],
 beer['brewers_tips'],
 beer['contributed_by'],
 beer['volume']['value'],
 ))) + '\n'
 for beer in beers
 ))
 cursor.copy_from(beers_string_iterator, 'staging_beers', sep='|', size=size)

La valeur par défaut pour la taille est 8192, qui est 2 ** 13 , nous conserverons donc les tailles en puissances de 2 :

>>> copy_string_iterator(connection, iter(beers), size=1024)
copy_string_iterator(size=1024)
Time 0.4536
Memory 0.0

>>> copy_string_iterator(connection, iter(beers), size=8192)
copy_string_iterator(size=8192)
Time 0.4596
Memory 0.0

>>> copy_string_iterator(connection, iter(beers), size=16384)
copy_string_iterator(size=16384)
Time 0.4649
Memory 0.0

>>> copy_string_iterator(connection, iter(beers), size=65536)
copy_string_iterator(size=65536)
Time 0.6171
Memory 0.0

Contrairement aux exemples précédents, il semble qu'il n'y ait pas de compromis entre la vitesse et la mémoire. Cela a du sens car cette méthode a été conçue pour ne consommer aucune mémoire. Cependant, nous obtenons un timing différent lors de la modification de la taille de la page. Pour notre ensemble de données, la valeur par défaut 8192 est le point idéal.

Résumé des résultats

Un résumé des résultats :