Датасеты 🤝 Arrow

Что такое Arrow?

Arrow позволяет быстро обрабатывать и перемещать большие объемы данных. Это специфический формат данных, который хранит их в столбчатой структуре памяти. Это обеспечивает несколько существенных преимуществ:

• Стандартный формат Arrow позволяет выполнять чтение без копирования, что практически полностью устраняет накладные расходы на сериализацию.
• Arrow не зависит от языка программирования, поэтому поддерживает различные языки.
• Arrow ориентирован на столбцы, поэтому запросы и обработка срезов или столбцов данных выполняются быстрее.
• Arrow позволяет передавать данные без копирования в стандартные инструменты машинного обучения, такие как NumPy, Pandas, PyTorch и TensorFlow.
• Arrow поддерживает множество, возможно вложенных, типов столбцов.

Отображение в память

🤗 Datasets использует Arrow для своей локальной системы кэширования. Это позволяет использовать на-дисковом кэше в качестве основы для наборов данных, который отображается в память (memory-mapped) для быстрого поиска. Такая архитектура позволяет работать с большими наборами данных на машинах с относительно небольшим объёмом памяти устройства.

Например, загрузка полного датасета английской Википедии занимает всего несколько МБ оперативной памяти:

>>> 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

Это возможно, поскольку данные Arrow фактически отображаются в память с диска (memory-mapped), а не загружаются в оперативную память. Memory-mapping позволяет обращаться к данным на диске и использует возможности виртуальной памяти для быстрого поиска.

Производительность

Итерация по датасету, отображённому в память с помощью Arrow, выполняется быстро. Итерация по Википедии на ноутбуке обеспечивает скорость 1–3 Гбит/с:

>>> 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