Apache Arrow :零内存读取DataFrame

本文翻译自Dejan Simic的文章《Apache Arrow: Read DataFrame With Zero Memory》
参考链接：Medium原文

上周我看到了 Wes McKinney 的一条推文，他是最出名的 pandas 库的创造者之一：

所以，当我看到他引用 William Gibson 的话时，我猜想一定有什么了不起的事情发生了。结果事情果然如我所料。

文章开头的推文是关于自然语言处理库 Hugging Face 的。该项目收集可用于模型训练和基准测试的数据集。其中一些数据集非常庞大。而在最初的推文中，Thomas Wolf 指出，通过特殊的文件格式，他和 Quentin Lhoest 现在能够在不到一分钟的时间内迭代 17GB 的数据，内存占用为 9MB！！

而我心里浮现的第一个问题就是：怎么做？ 这里是用了什么魔法吗？

而这就是 开头这篇 Twitter中 Wes McKinney 所谈论的未来！

由于网上没有具体实际的案例，因此决定写一篇介绍性的博客文章，其中包含迄今为止所学Arrow的动手示例。在本文末尾，你将找到所有来源的链接。

Any sufficiently advanced technology is indistinguishable from magic.
— Arthur C. Clarke, 3rd law

流畅的数据交换

第一条线索是 Thomas Wolf 提到 Apache Arrow。 Apache Arrow 是 Wes McKinney 发起的一个项目，旨在创建一个数据交换接口：

Apache Arrow 是内存数据的跨语言开发平台。它为平面和分层数据指定了一种与语言无关的标准化列式存储格式，为现代硬件上的高效分析操作而诞生。它还提供计算库和零拷贝流式消息传递和进程间通信。 [1]

这意味着什么？

在 Arrow 之前，任何应用程序或库之间交换数据的标准方式是以一种或另一种方式将其存储到磁盘。因此，如果 .NET Core 库想要将数据传递给 Python 进行数据分析，很可能有人会将数据写入文件（例如 csv、json、Parquet 等），然后用 Python 再次读取它。写入（序列化）和读取（反序列化）这两个步骤都是昂贵且缓慢的——数据集越大，完成每个步骤所需的时间就越长。

那是否存在一种通过握手和零复制直接交换数据的方式呢？举个栗子：.NET 将开始与 Python 聊天，指向内存中的一堆数据，然后就像：嘿，伙计，这堆东西现在是你的了。 然后Python 就可以直接访问这个数据，而不用把它从一个地方拖到另一个地方(例如先写入再读出)。那岂不是很棒？

而这，就是 Apache Arrow 的意义所在！

Parquet是Arrow实现的关键吗？

有了上面的介绍，这让我想知道——我该如何使用 Arrow？通过查看 Hugging Face 的源代码，了解到该项目使用 PyArrow 读取数据。自然而然地，我不自觉地将 PyArrow 与 Parquet 相联系了起来，Parquet 是一种高度压缩的列式存储格式。那么，Parquet 是 Arrow 交换数据的方式吗？ （剧透：不是）

按照传统地方式，数据以逐行方式存储在磁盘上。列式存储的诞生是为了分析大型数据集并有效地聚合它们。数据分析对数据的某一行不太感兴趣（例如，一笔客户交易、一份通话记录等），而是对其聚合（例如，客户花费的总金额、按地区划分的总通话分钟数等等）。

这导致了存储方向的改变：列式存储不是逐行存储，而是逐列排列数据。

而Parquet 就是一种基于列式存储的文件格式，它有两个主要优点 [4]：

• 高度可压缩性：.json 或 .csv 文件通常默认未压缩，而 Parquet 会压缩数据，因此可以节省大量磁盘空间。表 通常包含具有大量唯一值的列（high cardinality(这里理解为出现次数多且值不同)；例如用户 ID）和只有少数唯一值的列（low cardinality(虽然出现次数很多，但值都相同)；例如国家）。频率越低，压缩效果越好（可能）——下一节将详细介绍。

• 文件查询/过滤下推：在读入之前修剪不必要的数据。这可以缩短加载时间并优化资源消耗。如果需要从一千个表中获取两列，你不需要扫描所有行来获取这两个属性——你直接获取整个这两列。

压缩

为了更好地理解 Parquet 和 Arrow 之间的区别，我们需要绕道而行并获得一些关于压缩的知识。文件压缩本身就是一个巨大的主题。以下是一个简化的叙述，通过我自己对主题的理解过滤。这个题外话将有助于回答这两个问题：

• Parquet 如何将文件缩小到如此小的尺寸？
• Parquet 与 Arrow 有什么不同？

翻转硬币案例

这里可以想象一下你掷硬币十次并记录结果：

[正, 正, 正, 正, 反, 反, 反, 正, 反, 反]

如果我们要缩短这个结果并记录，你会怎样做？很有可能是 “4正3反1正2反”：

[4*正, 3*反, 正, 2*反]

这就是实际的压缩（所描述的算法称为游程编码 [8]）。我们倾向于自然地看到一种模式并缩写。压缩算法也可以做到这一点——只是具有更原始的计算能力和复杂的规则。但是这个例子应该足以帮助我们理解一个关键区别： .csv 文件采用字符的方式拼出每一条记录，而 Parquet 则是进行缩写（不会丢失任何信息）。

这个简单的例子也足以让我们直观地了解为什么压缩比会有很大差异。例如，如果排序顺序无关紧要，而你只是对 正 与 反 的总出现次数感兴趣，则可以先对列表进行排序，然后压缩版本如下所示：

[5*正, 5*反]

更进一步，如果我们在将数据集保存到 Parquet 之前先按所有列对数据集进行排序，则与未排序的文件相比，文件大小会更小。low cardinality越低(重复率越高)，压缩比越高。预计每列的压缩率将在排序顺序中进一步缩小。

Arrow 被压缩了吗？

基于以上分析，我们对为什么 Parquet 文件与未压缩的 .csv 文件相比如此之小有了一些认识。但这与Arrow有什么关系？

事实证明，这正是关键差异之一。 Parquet 以一种高效的方式存储在磁盘上。通过过滤器下推，您可以减少读入的数据量（即仅选择您实际需要的列）。但是当你想对数据进行操作时，你的计算机仍然需要将压缩后的信息解压缩并读入内存。 [2]

另一方面，Arrow是一种内存映射格式。在一篇博文中，Wes McKinney 将其总结如下：

Arrow 序列化设计提供了一个“数据头”，它描述了表中所有列的所有内存缓冲区的确切位置和大小。这意味着你可以内存映射巨大的、大于 RAM 的数据集，并在它们上就地评估 pandas 风格的算法，而无需像现在使用 pandas 那样将它们加载到内存中。你可以从 1 TB 表的中间读取 1 兆字节，而只需支付执行这些总计 1 兆字节的随机读取的成本

简而言之，应用程序可以直接对存储在磁盘上的数据集进行操作，而无需将其完全加载到内存中。正如最开始的那条推文所说。

实验：性能对比

现在让我们探索这些数据格式。作为示例数据集，我使用的是Palmer Station Penguin 数据集 。由于它只包含 350 行，因此我将其重新采样为 100 万行，以便性能差异变得更加明显：

## Read Palmer Station Penguin dataset from GitHub
import pandas as pd
df = pd.read_csv("https://raw.githubusercontent.com/allisonhorst/"
                 "palmerpenguins/47a3476d2147080e7ceccef4cf70105c808f2cbf/"
                 "data-raw/penguins_raw.csv")
                 # Increase dataset to 1m rows and reset index
df = df.sample(1_000_000, replace=True).reset_index(drop=True)


# Update sample number (0 to 999'999)
df["Sample Number"] = df.index
# Add some random variation to numeric columns
df[["Culmen Length (mm)", "Culmen Depth (mm)", 
    "Flipper Length (mm)", "Body Mass (g)"]] = df[["Culmen Length (mm)", "Culmen Depth (mm)", 
                                                   "Flipper Length (mm)", "Body Mass (g)"]] \
                                               + np.random.rand(df.shape[0], 4)

# Create dataframe where missing numeric values are filled with zero
df_nonan = df.copy()
df_nonan[["Culmen Length (mm)", "Culmen Depth (mm)", 
          "Flipper Length (mm)", "Body Mass (g)"]] = df[["Culmen Length (mm)", "Culmen Depth (mm)", 
                                                         "Flipper Length (mm)", "Body Mass (g)"]].fillna(0)

写入文件
下一步，我将以三种格式将文件写入磁盘：

• csv（包含缺失值）
• Parquet（包含缺失值）
• Arrow（包含 & 不包含缺失值）

请注意，在某些条件下，Arrow 只能在不分配内存（= 零拷贝）的情况下转换为 pandas。其中之一：必须没有 NaN 值。为了比较有和没有零复制的性能，我按照有无缺失数值编写了一次Arrow文件。

# Write to csv
df.to_csv("penguin-dataset.csv")

# Write to parquet
df.to_parquet("penguin-dataset.parquet")

# Write to Arrow
# Convert from pandas to Arrow
table = pa.Table.from_pandas(df)
# Write out to file
with pa.OSFile('penguin-dataset.arrow', 'wb') as sink:
    with pa.RecordBatchFileWriter(sink, table.schema) as writer:
        writer.write_table(table)

# Convert from no-NaN pandas to Arrow
table_nonan = pa.Table.from_pandas(df_nonan)
# Write out to file
with pa.OSFile('penguin-dataset-nonan.arrow', 'wb') as sink:
    with pa.RecordBatchFileWriter(sink, table_nonan.schema) as writer:
        writer.write_table(table_nonan)

生成的文件大小对比如下：

Parquet 生成的文件最小——尽管是随机序列（在写出文件之前没有进行排序），但它的压缩率高达 80%。Arrow 仅比 csv 略小。原因是 csv 甚至将数值存储为占用更多磁盘空间的字符串。在所有情况下，带有缺失值和不带有缺失值的文件之间的大小差异是微不足道的（<0.5 MB）。

读取时间

接下来是关键部分：读入性能。

• CSV
• Parquet
• 通过文件接口（OSFile()）读取Arrow
• 通过内存映射接口（memory_map()）读取Arrow(包含缺失值)
• 通过内存映射接口（memory_map()）读取Arrow(不包含缺失值)

# Read csv and calculate mean
%%timeit
pd.read_csv("penguin-dataset.csv")["Flipper Length (mm)"].mean()


# Read parquet and calculate mean
%%timeit
pd.read_parquet("penguin-dataset.parquet", columns=["Flipper Length (mm)"]).mean()


# Read Arrow using file API and calculate mean
%%timeit
with pa.OSFile('penguin-dataset.arrow', 'rb') as source:
    table = pa.ipc.open_file(source).read_all().column("Flipper Length (mm)")
result = table.to_pandas().mean()


# Read Arrow with memory-mapped API with missing values
%%timeit
source = pa.memory_map('penguin-dataset.arrow', 'r')
table = pa.ipc.RecordBatchFileReader(source).read_all().column("Flipper Length (mm)")
result = table.to_pandas().mean()


# Read Arrow with memory-mapped API without missing values (zero-copy)
%%timeit
source = pa.memory_map('penguin-dataset-nonan.arrow', 'r')
table = pa.ipc.RecordBatchFileReader(source).read_all().column("Flipper Length (mm)")
result = table.to_pandas().mean()

对不同实现方法进行计时，整体的对比结果如下：

不出所料， csv 是最慢的选择。它需要读取 200MB，解析文本，丢弃除长度以外的所有列，然后计算平均值。

Parquet 的速度快了约 60 倍，因为不需要解析整个文件——只读入所需的列。

包含缺失值的 Arrow 比 Parquet 快约 3 倍，比 csv 快约 200 倍。与 Parquet 一样，Arrow 可以将自身限制为仅读取指定的列。而它比 Parquet 的原因是它不需要解压缩列(Parquet存储小，但处理时需要解压缩)。

请注意，使用和不使用零复制读取内存映射 Arrow 文件之间的差异意味着性能又提高了约 3 倍（即，零复制总共比 csv 快约 600 倍，比 Parquet 快约 9 倍）。

令人惊讶的是：通过文件接口（OSFile()）读取Arrow 甚至比 Parquet 还要慢。这是怎么回事？

内存消耗

为了回答这个问题，让我们看看内存消耗。如果我们读取单个列，每个文件消耗多少 RAM？

# Measure initial memory consumption
memory_init = psutil.Process(os.getpid()).memory_info().rss >> 20


# Read csv
col_csv = pd.read_csv("penguin-dataset.csv")["Flipper Length (mm)"]
memory_post_csv = psutil.Process(os.getpid()).memory_info().rss >> 20


# Read parquet
col_parquet = pd.read_parquet("penguin-dataset.parquet", columns=["Flipper Length (mm)"])
memory_post_parquet = psutil.Process(os.getpid()).memory_info().rss >> 20


# Read Arrow using file API
with pa.OSFile('penguin-dataset.arrow', 'rb') as source:
    col_arrow_file = pa.ipc.open_file(source).read_all().column("Flipper Length (mm)").to_pandas()
memory_post_arrowos = psutil.Process(os.getpid()).memory_info().rss >> 20


# Read Arrow with memory-mapped API with missing values
source = pa.memory_map('penguin-dataset.arrow', 'r')
table_mmap = pa.ipc.RecordBatchFileReader(source).read_all().column("Flipper Length (mm)")
col_arrow_mapped = table_mmap.to_pandas()
memory_post_arrowmmap = psutil.Process(os.getpid()).memory_info().rss >> 20


# Read Arrow with memory-mapped API without missing values (zero-copy)
source = pa.memory_map('penguin-dataset-nonan.arrow', 'r')
table_mmap_zc = pa.ipc.RecordBatchFileReader(source).read_all().column("Flipper Length (mm)")
col_arrow_mapped_zc = table_mmap_zc.to_pandas()
memory_post_arrowmmap_zc = psutil.Process(os.getpid()).memory_info().rss >> 20


# Display memory consumption
print(f"csv memory consumption: {memory_post_csv - memory_init}\n"
      f"Parquet memory consumption: {memory_post_parquet - memory_post_csv}\n"
      f"Arrow file memory consumption: {memory_post_arrowos - memory_post_parquet}\n"
      f"Arrow mapped (no zero-copy) memory consumption: {memory_post_arrowmmap - memory_post_arrowos}\n"
      f"Arrow mapped (zero-copy) memory consumption: {memory_post_arrowmmap_zc - memory_post_arrowmmap}\n")

实验结果如下：

值得注意的是：通过文件接口（OSFile()）读入Arrow 消耗了 189 MB——这几乎是整个文件的大小（即使我们只读取单个列？！）。而答案可以在Arrow文档中找到：

“[…] OSFile allocates new memory on each read, like Python file objects.” [3]
OSFile 在每次读取时分配新内存，就像 Python 读取文件对象一样。” [3]

通过使用 OSFile，整个文件首先被读入内存。现在很明显为什么这个操作比 Parquet 慢并且消耗最多的内存！

然而，通过使用内存映射函数和填充 NaN 值，pandas DataFrame 是直接在存储的 Arrow 文件之上创建的。无复制：0 MB RAM！毫不奇怪，这是最快的选择。

整个的性能对比notebook如下Arrow性能对比

结论

对我来说，关于 Arrow 有很多值得学习的地方。到目前为止我学到的是：鱼和熊掌不可兼得。 下面这两个方式中需要进行权衡：

• 优化磁盘空间/磁盘上的长期存储 → Parquet
• 优化数据交换和快速检索 → Arrow

与 csv 相比，Parquet 和 Arrow 的性能提升非常显着。将 Arrow 存储到磁盘时，它比 Parquet 消耗更多的存储空间。然而，Arrow 在阅读性能方面胜过 Parquet——无论是在时间还是内存消耗方面。所提供的示例（计算一列/读取列的平均值）只是表面上的问题——我的期望是，随着更复杂的查询和更大的数据集，Arrow 会更加闪耀。

只要使用内存映射功能读取 Arrow，读取性能就令人难以置信。最好的情况是数据集没有缺失值/NaN。然后 PyArrow 可以发挥它的魔力，让你在桌子上进行操作，几乎不消耗任何内存。

最后，感慨一下：未来已来，加油吧！

[1]: Apache Arrow,Landing Page(2020), Apache Arrow 网站
[2]: Apache Arrow, FAQ (2020), Apache Arrow 网站
[3]: Apache Arrow，磁盘和内存映射文件 (2020)，Apache Arrow Python 文档
[4]: J. LeDem, Apache Arrow and Apache Parquet: Why We Needed Different Projects for Columnar Data, On Disk and In-Memory(2017), KDnuggets
[5] J. LeDem, The Columnar Roadmap: Apache Parquet and Apache Arrow (2017), Dremio
[6] W. McKinney, Apache Arrow and the “10 Things I Hate About pandas” (2017), 博客
[7] W. McKinney, Some comments to Daniel Abadi’s blog about Apache Arrow (2017), 博客
[8] 维基百科, Run-length encoding (2020), 维基百科