Jeux de données 🤝 Arrow

Qu'est-ce qu'Arrow ?

Arrow permet de traiter et de déplacer de grandes quantités de données rapidement. C'est un format de données spécifique qui stocke les données dans une disposition mémoire en colonnes. Cela offre plusieurs avantages majeurs :

• Le format standard d’Arrow permet des lectures sans copie, ce qui élimine pratiquement toute la surcharge de sérialisation.
• Arrow est indépendant du langage, il prend donc en charge différents langages de programmation.
• Arrow est orienté colonnes, ce qui le rend plus rapide pour interroger et traiter des tranches ou des colonnes de données.
• Arrow permet des transferts sans copie vers des outils standard d'apprentissage automatique tels que NumPy, Pandas, PyTorch et TensorFlow.
• Arrow prend en charge de nombreux types de colonnes, éventuellement imbriqués.

Mappage mémoire

🤗 Datasets utilise Arrow pour son système de cache local. Il permet aux jeux de données d’être adossés à un cache sur disque, mappé en mémoire pour une recherche rapide. Cette architecture permet d’utiliser de grands jeux de données sur des machines dotées d’une mémoire de périphérique relativement limitée.

Par exemple, charger l'ensemble de données complet de Wikipédia en anglais ne nécessite que quelques Mo de RAM :

>>> import os; import psutil; import timeit
>>> from datasets import load_dataset

# Process.memory_info is expressed in bytes, so convert to megabytes 
>>> mem_before = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
>>> wiki = load_dataset("wikimedia/wikipedia", "20220301.en", split="train")
>>> mem_after = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)

>>> print(f"RAM memory used: {(mem_after - mem_before)} MB")
RAM memory used: 50 MB

Cela est possible car les données Arrow sont en réalité mappées en mémoire depuis le disque, et non chargées en mémoire. Le mappage mémoire permet d’accéder aux données sur le disque et exploite les capacités de mémoire virtuelle pour des recherches rapides.

Performance

Itérer sur un jeu de données en mémoire mappée à l’aide d’Arrow est rapide. Itérer sur Wikipédia sur un ordinateur portable donne des débits de 1 à 3 Gbit/s :

>>> s = """batch_size = 1000
... for batch in wiki.iter(batch_size):
...     ...
... """

>>> elapsed_time = timeit.timeit(stmt=s, number=1, globals=globals())
>>> print(f"Time to iterate over the {wiki.dataset_size >> 30} GB dataset: {elapsed_time:.1f} sec, "
...       f"ie. {float(wiki.dataset_size >> 27)/elapsed_time:.1f} Gb/s")
Time to iterate over the 18 GB dataset: 31.8 sec, ie. 4.8 Gb/s