Datasets 🤝 Arrow

¿Qué es Arrow?

Arrow permite procesar y mover grandes cantidades de datos rápidamente. Es un formato de datos específico que almacena los datos en un diseño de memoria columnar. Esto ofrece varias ventajas significativas:

• El formato estándar de Arrow permite lecturas sin copia, lo que elimina prácticamente toda la sobrecarga de serialización.
• Arrow es independiente del lenguaje, por lo que admite diferentes lenguajes de programación.
• Arrow está orientado a columnas, por lo que es más rápido al consultar y procesar fragmentos o columnas de datos.
• Arrow permite transferencias sin copia a herramientas estándar de aprendizaje automático como NumPy, Pandas, PyTorch y TensorFlow.
• Arrow admite muchos tipos de columnas, posiblemente anidados.

Mapeo de memoria

🤗 Datasets utiliza Arrow para su sistema de caché local. Esto permite que los conjuntos de datos se respalden en una caché en disco, que se mapea en memoria para una búsqueda rápida. Esta arquitectura permite utilizar conjuntos de datos grandes en máquinas con una memoria de dispositivo relativamente pequeña.

Por ejemplo, cargar el conjunto de datos completo de Wikipedia en inglés solo ocupa unos pocos MB 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

Esto es posible porque los datos de Arrow están realmente mapeados en memoria desde el disco, y no cargados en la memoria. El mapeo en memoria permite acceder a los datos en el disco y aprovecha las capacidades de memoria virtual para búsquedas rápidas.

Rendimiento

Iterar sobre un conjunto de datos mapeado en memoria usando Arrow es rápido. Iterar sobre Wikipedia en un portátil ofrece velocidades de 1 a 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