数据集文档

数据集 🤝 Arrow

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验

开始使用

数据集 🤝 Arrow

什么是 Arrow?

Arrow 使得能够快速处理和移动大量数据。它是一种特定的数据格式,以列式内存布局存储数据。这提供了几个显著的优势

  • Arrow 的标准格式允许 零拷贝读取,从而几乎消除了所有序列化开销。
  • Arrow 与语言无关,因此它支持不同的编程语言。
  • Arrow 是面向列的,因此在查询和处理数据切片或列方面速度更快。
  • Arrow 允许无复制地传递到标准机器学习工具,例如 NumPy、Pandas、PyTorch 和 TensorFlow。
  • Arrow 支持许多(可能嵌套的)列类型。

内存映射

🤗 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("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 上更新