数据集 🤝 Arrow

什么是 Arrow？

Arrow 能够快速处理和传输大量数据。它是一种特定的数据格式，以列式内存布局存储数据。这带来了多项显著优势：

Arrow 的标准格式支持零拷贝读取，从而几乎完全消除了序列化开销。
Arrow 与语言无关，因此支持不同的编程语言。
Arrow 采用列式存储，因此在查询和处理数据切片或列时速度更快。
Arrow 支持向 NumPy、Pandas、PyTorch 和 TensorFlow 等标准机器学习工具进行零拷贝数据传递。
Arrow 支持许多列类型，这些类型可能相互嵌套。

内存映射

🤗 Datasets 使用 Arrow 作为其本地缓存系统。它允许数据集由磁盘缓存提供支持，并通过内存映射实现快速查找。这种架构使得在设备内存相对较小的机器上也能使用大型数据集。

例如，加载完整的英文维基百科数据集仅需占用几 MB 的内存：

>>> 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）的，而非加载到内存中。内存映射允许访问磁盘上的数据，并利用虚拟内存机制实现快速查找。

性能

使用 Arrow 遍历内存映射数据集的速度很快。在笔记本电脑上遍历维基百科，速度可达 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

Update on GitHub