Datasets 🤝 Arrow

什么是 Arrow?

Arrow 能够快速处理和移动大量数据。它是一种特定的数据格式，将数据存储在列式内存布局中。这带来了几个显著的优势：

Arrow 的标准格式允许零拷贝读取 (zero-copy reads)，这几乎消除了所有的序列化开销。
Arrow 与语言无关，因此支持不同的编程语言。
Arrow 是面向列的，因此在查询和处理数据的切片或列时速度更快。
Arrow 允许将数据在无需拷贝的情况下直接传递给标准的机器学习工具，如 NumPy、Pandas、PyTorch 和 TensorFlow。
Arrow 支持许多列类型，包括可能的嵌套类型。

内存映射 (Memory-mapping)

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

例如，加载完整的英文维基百科数据集仅需几 MB 的 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

这是因为 Arrow 数据实际上是从磁盘内存映射的，而不是加载到内存中。内存映射允许访问磁盘上的数据，并利用虚拟内存能力实现快速查找。

性能

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

在 GitHub 上更新