Datasets 🤝 Arrow

O que é Arrow?

Arrow permite que grandes volumes de dados sejam processados e movidos rapidamente. É um formato de dados específico que armazena dados em um layout de memória colunar. Isso oferece várias vantagens significativas:

• O formato padrão do Arrow permite leituras sem cópia, o que elimina virtualmente todo o sobrecarga de serialização.
• O Arrow é agnóstico de linguagem, portanto, suporta diferentes linguagens de programação.
• O Arrow é orientado a colunas, o que o torna mais rápido ao consultar e processar fatias ou colunas de dados.
• O Arrow permite a transferência de dados sem cópia para ferramentas padrão de aprendizado de máquina, como NumPy, Pandas, PyTorch e TensorFlow.
• O Arrow suporta muitos tipos de colunas, possivelmente aninhados.

Mapeamento de memória

🤗 O Datasets usa Arrow para seu sistema de cache local. Isso permite que os conjuntos de dados sejam suportados por um cache em disco, que é mapeado na memória para busca rápida. Essa arquitetura permite o uso de grandes conjuntos de dados em máquinas com memória de dispositivo relativamente pequena.

Por exemplo, carregar o conjunto de dados completo da Wikipédia em inglês ocupa apenas alguns 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

Isso é possível porque os dados do Arrow são, na verdade, mapeados na memória a partir do disco, e não carregados na memória. O mapeamento de memória permite o acesso a dados no disco e aproveita as capacidades de memória virtual para consultas rápidas.

Desempenho

Iterar sobre um dataset mapeado em memória usando Arrow é rápido. Iterar sobre a Wikipedia em um laptop resulta em velocidades 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