Parser des fichiers XML de plusieurs gigaoctets avec Python

TLDR; → Voir le tableau récapitualtif.

Pour le projet eBook Reader Dictionaries, il faut parser des fichiers XML de plusieurs gigaoctets, et ça pose quelques problèmes de performance.

En l'occurrence le fichier est un export du Wiktionnaire (fichier pages-20200420.xml) qui pèse quelques 4,2 Gio.
Il contient 4 107 154 mots, chacun encapsulé dans ce shéma :

<page>
    <title>$word</title>
    <ns>int</ns>
    <id>str</id>
    <revision>
        <id>str</id>
        <parentid>str</parentid>
        <timestamp>str</timestamp>
        <contributor>
            <username>str</username>
            <id>str</id>
        </contributor>
        <comment>str</comment>
        <model>wikitext</model>
        <format>text/x-wiki</format>
        <text size="int" xml:space="preserve">$wikicode</text>
    </revision>
</page>

Tenter de charger et parser un tel fichier va mettre à mal votre machine. J'ai 32 Go de RAM et j'ai réussi à saturer ma Swap... Donc voyons comment charger et parser ce fichier par itérations.

Les soucis de performance se posent à 2 moments :

Au parsage pur et dur du document.
Au traitement des éléments du document.

Cet article traite du point n°1 en priorité, bien que certaines méthodes soient influencées par le point n°2 (suivant le nombre, la complexité et les données contenues dans chaque élément).

Voici le décorateur qui servira de chronomètre :

from datetime import timedelta
from time import monotonic


def timer(func):
    def inner(file):
        start = monotonic()
        try:
            func(file)
        finally:
            end = monotonic()
            delta = timedelta(seconds=end - start)
            print(f"{func.__name__}(): {delta}", flush=True)
    return inner


@timer
def function(file: str) -> None:
    pass

Idée 1 : xmltodict

Je m'étais d'abord orienté vers xmltodict pour sa simplicité : tu lui fournis un fichier XML, aussi gros soit-il, et tu récupères un dict de chaque élément .

Ces imports seront nécessaires :

from typing import Any, Dict
import xmltodict

Le module travaille avec une fonction callback. Le traitement de chaque élément se fera dans cette fonction.

def handle_element(_, page: Dict[str, Any]) -> bool:
    """
    Callback passed to xmltodict.parse().
    The function must return True or the parser will raise ParsingInterrupted
    (https://github.com/martinblech/xmltodict/blob/d6a8377/xmltodict.py#L227-L230).
    """
    # (https://github.com/BoboTiG/ebook-reader-dict/blob/2a275c6/scripts/get.py#L223)
    return True

Il suffit ensuite de démarrer le bousin :

with file.open("rb") as fh:
    xmltodict.parse(fh, item_depth=2, item_callback=handle_element)

Le code est épuré pour ne garder que la logique de parsage. Le traitement de chaque élément n'est pas pertinent et se trouve dans le code sur GitHub.

Voyons ce que ça donne :

$ python func1.py
Time: 0h:07m:27s

Idée 2 : ElementTree

7 minutes me paraissaient (beaucoup) trop long, j'ai donc ouvert une issue et passé un peu de temps dessus.
Je me suis penché vers ElementTree de la bibliothèque standard. Étant en Python 3.7, le module utilise déjà une implémentation écrite en C (le bon vieux temps de cElementTree est révolu).

Ces imports seront nécessaires :

import xml.etree.ElementTree as ET
from typing import Generator, Tuple

Le code est ensuite découpé :

Une fonction qui parse le document XML et yield les éléments.
Un autre fonction pour traiter chaque élément.

def xml_iter_parse(file: str) -> Generator[ET.Element, None, None]:
    """Efficient XML parsing for big files.
    Elements are yield'ed when they meet the "page" tag.
    """
    start_tag = None
    doc = ET.iterparse(file, events=("start", "end"))
    _, root = next(doc)

    # This is specific to Wiktionary dumps
    tag_page = "{http://www.mediawiki.org/xml/export-0.10/}page"

    for event, element in doc:
        if start_tag is None and event == "start" and element.tag == tag_page:
            start_tag = element.tag
        elif start_tag is not None and event == "end" and element.tag == start_tag:
            yield element
            start_tag = None
            root.clear()  # Keep memory low

def xml_parse_element(element: ET.Element) -> Tuple[str, str, str]:
    """Parse the *element* to retrieve the word, the current revision and definitions."""
    # (https://github.com/BoboTiG/ebook-reader-dict/blob/50ff3b2/scripts/get.py#L292)
    return "", "", ""

La fonction xml_iter_parse() est spécifique aux exports de Wiktionnaire, mais elle peut être rendue générique en supprimant les 2 occurrences de tag_page ; ainsi, tous les éléments seront yield'és.

Démarrons la bête :

for element in xml_iter_parse(file):
    word, rev, code = xml_parse_element(element)

Voyons ce que ça donne :

$ python func2.py
Time: 0h:02m:08s

Wow, un joli gain comparé à la version précédente !

Idée 3 : SAX

Suite à la suggestion de Habib, j'ai testé le module xml.sax.

Ces imports seront nécessaires :

import xml.sax
import xml.sax.handler
from typing import Dict

Le code est ensuite découpé :

Une classe qui représentera les éléments du document XML. C'est intéressant car c'est ici que le filtrage des balises peut se faire.
La logique de l'instanciation du parseur.

class PageHandler(xml.sax.handler.ContentHandler):
    """XML content handler passed to the SAX parser."""

    __slots__ = ("_current_tag", "_is_page", "code", "word", "words")

    def __init__(self) -> None:
        # Will contain the Wikicode
        self.code = ""
        # Will contain the word
        self.word = ""

        self._current_tag = ""
        self._is_page = False

    def reset(self) -> None:
        """Reset attributes."""
        self.code = ""
        self.word = ""
        self._is_page = False

    def startElement(self, tag: str, attributes: Dict[str, str]) -> None:
        """Signals the start of an element."""
        if tag == "page":
            self._is_page = True
        elif tag in ("title", "text"):
            self._current_tag = tag

    def endElement(self, tag: str) -> None:
        """Signals the end of an element."""
        if self._current_tag == "text":
            if self.word:
                # The parsing is complete!
                self.process()
            self.reset()
        self._current_tag = ""

    def process(self) -> None:
        """Process the element."""
        pass

    def characters(self, content: str) -> None:
        """Receive notification of character data."""
        if not self._is_page:
            # We are not traiting a <page> element, not interesting
            pass
        elif self._current_tag == "title":
           self.word = content
        elif self._current_tag == "text":
            # Aggregate all the Wikicode
            self.code += content

C'est assez spécifique au projet, mais c'est plus simple à comprendre avec du vrai code.

Maintenant, créons le parseur :

parser = xml.sax.make_parser()

# All element names, prefixes, attribute names, Namespace URIs, and local names
# are interned using the built-in intern function.
parser.setFeature(xml.sax.handler.feature_string_interning, 1)

handler = PageHandler()
parser.setContentHandler(handler)
parser.parse(file)

Voyons ce que ça donne :

$ python func3.py
Time: 0h:05m:46s

C'est un peu mieux que xmltodict, mais plus long que ElementTree.

J'espèrais gratter quelques minutes grâce au filtrage des balises. J'imagine que le va-et-vient entre le code C <-> Python est le goulot d'étranglement ici.

Résumé

Tableau récapitulatif des solutions testées :

Module	Temps de parsage	Ratio
xmltodict	0h:07m:27s	3,49
ElementTree	0h:02m:08s	1
SAX	0h:05m:46s	2,7

Pour la Suite

J'ai ouvert ces issues pour pousser l'optimisation un peu plus loin :

Smaller Elements and less work done in useless tags pour réduire la taille et le nombre des objets crées pour chaque élément (2020-05-06 : cf. l'idée #3, pas terrible finalement).
Distributing the workload pour enfin partager le travail sur plusieurs processus ou cœurs.
Une autre idée, peut-être ? Le code du benchmark se trouve ici : bench-xml-parser.py.

Au final, avec ElementTree le temps est réellement intéressant. Le gros du travail qui reste à gérer, c'est le traitement des éléments, là on atteint des sommets rapidement.

Historique

2020-05-06 : Étoffement des explications, ajout de l'idée n°3 (xml.sax) et du tableau récapitulatif.

Blog

Luma

Liens

GitHub

Parser des fichiers XML de plusieurs gigaoctets avec Python

Idée 1 : xmltodict

Idée 2 : ElementTree

Idée 3 : SAX

Résumé

Pour la Suite

Historique