Et justement, en relisant un célèbre post de Guido sur la réponse à la blague “comment trier un million d’entiers avec 2M de Ram”, j’ai réalisé 2 choses:

• Le contenu de l’article est génial.
• Le contenu de l’article est incompréhensible.

Or c’est un peu ma raison d’être, si vous voulez, de prendre les trucs cools mais imbitables et les rendre utilisables.

Aujourd’hui, donc, on va trier un CSV de 5Go en Python. Ça parle plus qu’un fichier de nombres, et puis ça évite d’expliquer le module array qui est juste une optimisation.

J’ai pris mes petites mimines, et j’ai pondu un CSV qui contient 63Mo de:

A,01/02/2016,2
A,13/07/2011,1
B,24/01/1996,3
C,30/12/1999,1
D,13/07/2011,3
D,01/02/2016,5
E,24/01/1996,4
F,30/12/1999,1
G,13/07/2011,4
H,01/02/2016,4
I,01/02/2016,5
I,13/07/2011,2
A,01/02/2016,2
A,13/07/2011,1

En copier/coller.

Puis j’ai lancé un script pour dupliquer le bébé et jusqu’à atteindre 5,9 Go de taille:

with open('data.csv') as infile:
    with open('data2.csv', 'w') as outfile:
        for x in range(100):
            outfile.write(infile.read())
            infile.seek(0)

Si jamais vous doutiez que je vous aime…

395366400 lignes. Jusqu’ici tout va bien.

Maintenant, supposons qu’on veuille trier sur la date. Si vos souvenirs en Python sont exacts (ou si vous avez lu notre super article sur Ordonner en Python), vous savez que la solution naïve est de loader le fichier cash pistache, puis d’appeler dessus sorted() avec le paramètre key.

D’abord, il faut choisir le callback à passer à key, c’est à dire la fonction qui va être exécutée pour chaque ligne pour extraire la date de la ligne et permettre ainsi à sorted() de comparer avec les autres dates.

>>> str_date = "A,01/02/2016,2".split(',')[1] # récupère la date uniquement
>>> str_date
'01/02/2016'
>>> ''.join(str_date.split('/')[::-1]) # on inverse la date avoir une valeur ordonnable
'20160201'

On en fait une fonction:

 
def extract_date(ligne):
    str_date = ligne.split(',')[1]
    return ''.join(str_date.split('/')[::-1])

Ensuite on a juste à ouvrir le fichier, trier les lignes, et sauvegarder tout ça dans l’autre fichier:

with open('data2.csv') as infile:
    with open('sorted_data.csv', 'w') as outfile:
        # on fait le tri
        sorted_lines = sorted(infile, key=extract_date)
        # on ecrit les lignes triées dans le nouveau fichier
        outfile.writelines(sorted_lines)

Easy money, double poney.

Enfin avec un fichier de quelques Mo. Parce que si vous faites ça dans un fichier de 5,9 Go, votre RAM va vomir comme une pom pom girl anorexique.

sorted() sur un disque complet, illustré

Comment résoudre ce problème ?

Et bien en faisant tout pareil, mais avec des petits morceaux !

import heapq

from tempfile import TemporaryFile
from itertools import islice

# On garde notre fonction key
def extract_date(ligne):
    str_date = ligne.split(',')[1]
    return ''.join(str_date.split('/')[::-1])

# Liste de tous les fichiers temporaires avec les lignes triées
sorted_temp_files = []

with open('data2.csv') as infile:
    progress = 0
    while True:
        # On lit seulement 3000000 lignes sur 395366400 à chaque tour de boucle
        lines = list(islice(infile, 3000000))

        if not lines:  # plus de ligne ? On sort
            break

        # On affiche où on en est
        print("{:.2f}%".format(progress))
        progress += (3000000 / 395366400 * 100)

        # On tri les lignes, comme avec sorted() mais sur place. 
        # Ça évite de dupliquer les données en mémoire.
        lines.sort(key=extract_date)

        # On crée un fichier temporaire qui va contenir les 3000000 lignes
        # triées
        f = TemporaryFile(mode="r+")
        f.writelines(lines)

        # On rembobine le fichier pour pouvoir relire le contenu plus tard
        f.seek(0)

        # On balance le fichier dans la liste des fichiers triés
        # L'objet fichier hein. Pas le chemin du fichier. C'est pour ça qu'on
        # a fait .seek(0) avant. On va en avoir besoin plus bas.
        sorted_temp_files.append(f)

    # Toute la magie se fait là.
    # On utilise heapq.merge(), qui prend en paramètre plein de trucs qu'on 
    # vient de trier, et permet de se balader dedans comme si c'était un seul
    # gros truc trié, mais sans tout charger en mémoire.
    # En gros il regarde les premières valeurs de tous les itérables, les compares,
    # puis fait retourne un générateur qui yield tout dans l'ordre 
    with open('sorted_data.csv', 'w') as outfile:
        for ligne in heapq.merge(*sorted_temp_files, key=extract_date):
            outfile.write(ligne)

Au lieu de charger les 5 Go en mémoire, on plafonne à 400 Mo. Une raison de plus d’apprendre à utiliser les générateurs.

Sametmax, c’est de la bombe

Alors évidemment, c’est long. Y a genre 40 minutes de traitement. Moins si on utilise pypy et qu’on accepte de up la conso RAM. On code pas en Rust :)

Si vous avez besoin de perfs et que Python reste votre outil de prédilections, 3 solutions restent à votre disposition:

– Prétraiter le fichier en lui rajoutant une première colonne qui contient un timestamp. Ca c’est super rapide à trier.
– Utiliser un truc de numpy ou pandas comme np.memmap().sort(kind=’merge’). C’est du C, donc ça speed.
– Acheter de la ram et tout trier en mémoire avec la solution 1.

EDIT:

Un lecteur m’a alpagué sur twitter pour me dire qu’on peut faire ça plus vite et plus facilement avec sort. Oui, évidement.

L’exercice est académique, il est simplifié pour vous permettre de comprendre heapq.merge(). En réalité, on aura vraiment besoin de ça que pour un cas complexe. Exemple, vous lisez un flux de log une socket et vous checkez les adresses IP d’abord pour filtrer celles de la liste noire, puis pour les faire matcher une zone géographique et vous triez tout ça et vous le passez à la suite du pipeline. Et le tout doit marcher sous Linux et Windows Server avec juste la lib standard parce que diab est de mauvais poil.

Évidement que si vous voulez juste trier un CSV vous n’allez pas coder un script python de 30 lignes.

Ca aurait été sur ce blog, le mec se serait pris un tampon “drozophiliafucker” dans la gueule et ça aurait été plié. Mais quand on est BDFL et, à l’époque, employé chez Google, on est obligé de donner une réponse un peu sérieuse.

Si vous avez suivi le lien précédent, vous verrez que sa réponse implique un obscure module appelé heapq. Si vous allez sur la doc du-dit module, vous verrez qu’elle implique une obscure explication innomable et vous aller retourner à la vidéo de porn que vous avez laissé bufferiser en haute résolution afin de pouvoir voir les grains de beauté qui longent le pourtour anal de la principale protagoniste. Respirons.

Il y a une explication rationnelle à tout cela

heapq est un algorithme qui organise une liste sous forme d’arbre binaire. Vous voyez c’était simple non ?

Nan je déconne, je vais pas vous laisser avec ça quand même.

Le module heapq permet d’utiliser le type “list” pour y ajouter ou retirer les éléments de telle sorte qu’ils soient toujours dans l’ordre.

Sous le capot, ça marche effectivement avec un arbre binaire, mais on s’en branle. Tout ce qu’il faut comprendre c’est:

>>> from heapq import heappush, heappop
>>> l = []
>>> heappush(l, 69)
>>> heappush(l, 42)
>>> heappush(l, 2048)
>>> heappush(l, -273.15)
>>> l # la liste est ordonnée en arbre binaire...
[-273.15, 42, 2048, 69]
>>> for x in xrange(len(l)): # et on depop pour itérer dans l'odre
    print heappop(l)
...     
-273.15
42
69
2048

Donc on a une liste, on lui met des éléments dans la tronche dans n’importe quel ordre, et bam, on peut itérer dans le bon ordre sans avoir à rien faire. Et cette insertion est assez rapide (O(lg n) pour les tatillons). Le parcours l’est également (de l’ordre de O(n log n)).

Et donc c’est très pratique pour trouver les x plus petits éléments (ou les plus grands), implémenter des queues de priorité, etc.

Exemple de queue de priorité, (courageusement volé d’ici):

import heapq

class PriorityQueue(list):

    def __init__(self, data):
        super(Heap, self).__init__()
        for i, x in enumerate(data):
            self.push(i, x)
   
    def push(self, priority, item):
        """
            On push en rajoute une priorité
        """
        heapq.heappush(self, (priority, item))
   
    def pop(self):
        """
            On pop en retirant la proprité
        """
        return heapq.heappop(self)[1]
   
    def __len__(self):
        return len(self)
   
    def __iter__(self):
        """
            Comme on a une méthode next(), on peut se retourner soi-même.
            Ainsi la boucle for appelera next() automatiquement. 
        """
        return self
   
    def next(self):
        """ 
           On depop la liste du plus petit au plus grand.
        """
        try:
            return self.pop()
        except IndexError:
            raise StopIteration

>>> l = PriorityQueue(("azerty"))
>>> l
>>> l.push(100, 'après')
>>> l.push(-1, 'avant')
>>> l.push(5, 'pendant')
>>> for x in l:
...     print x
...     
avant
a
z
e
r
t
pendant
y
après

Et le plus beau, c’est qu’on peut prendre plusieurs itérables ordonnés, et utiliser heapq.merge pour obtenir un générateur (qui ne charge pas tout en mémoire d’un coup) qui va permettre d’iterer de manière ordonnée sur tous les éléments.

>>> import heapq
>>> l1 = sorted([random.randint(0, 1000) for x in xrange(5)])
>>> l2 = sorted([random.randint(0, 1000) for x in xrange(5)])
>>> l3 = sorted([random.randint(0, 1000) for x in xrange(5)])
>>> list(heapq.merge(l1, l2, l3))
[52, 59, 60, 171, 174, 262, 336, 402, 435, 487, 557, 645, 899, 949, 996]

Notez que ce n’est pas du tout la même chose que de concaténer les listes:

>>> l1 + l2 + l3
[59, 174, 336, 487, 996, 52, 171, 557, 645, 949, 60, 262, 402, 435, 899]

Car des élements de l2 peuvent être inférieurs à ceux de l1. heap.merge nous ordonne tout bien correctement. C’est l’équivalent de sorted(l1 + l2 + l3), sauf que ça ne charge pas tout en mémoire:

>>> heapq.merge(l1, l2, l3)

Alors que sorted(), charge tout en mémoire:

>>> sorted(l1 + l2 + l3)
[52, 59, 60, 171, 174, 262, 336, 402, 435, 487, 557, 645, 899, 949, 996]

Bien entendu, ça marche avec des listes créées avec
heapq, ainsi:

>>> l1 = []

>>> for x in xrange(5): heappush(l1, random.randint(0, 1000))

>>> l2 = []

>>> for x in xrange(5): heappush(l2, random.randint(0, 1000))

>>> l3 = []

>>> for x in xrange(5): heappush(l3, random.randint(0, 1000))

>>> list(heapq.merge(l1, l2, l3))

[31, 40, 133, 360, 504, 508, 513, 679, 645, 792, 838, 413, 765, 886, 924]

(Grosse connasserie de ma part, faites comme si vous aviez rien vu.)

Quand on a de grosses quantités de données à trier, c’est très pratique, car l’effort de tri est répartie à chaque insertion, et à chaque itération pendant le parcours, pas concentré sur de gros appels de sorted().

On peut aussi récupérer des trucs du genre, les n plus petits / grands éléments sans tout coder à la main:

>>> l = [random.randint(0, 1000) for x in xrange(100)]
>>> heapq.nsmallest(4, l)
[0, 2, 4, 7]
>>> heapq.nlargest(3, l)
[999, 996, 983]

Et c’est beaucoup plus efficace que de le faire soi-même.

J’en profite pour rappeler au passage que tous les objets en Python sont ordonnables:

>>> (1, 2) > (2, 1)
False
>>> (1, 2) < (2, 1)
True
>>> "a" > "b"
False
>>> "a" < "b"
True

Et qu'en définissant les méthodes __eq__, __lt__, __gt__, etc., on peut donner un moyen de comparer n'importe quelle classe.

Bref, pour tous les besoins de classement, de priorisations, d'ordonnancement, de notion de plus petits, de plus grands, etc. qui concernent un gros jeu de données, heapq est le module qu'il vous faut.

Et là je percute que ça fait quelque temps déjà que je fais des articles élitistes pour des uses cases de 1% de la population. Donc les prochains tutos concerneront des trucs plus terre à terre. Parce que, je fais le malin là, mais je savais pas à quoi servait heapq il y a 3 semaines.

xoxo les filles