使用 Apache Arrow 进行内存分析——格式与内存处理

971 阅读49分钟

我一直在称赞 Apache Arrow 作为表格数据交换技术的优势,但它与人们常用的数据传输技术相比表现如何呢?对于你的应用程序接口(API),什么时候使用一种技术比另一种更合适呢?要回答这些问题,你需要了解不同技术如何利用内存。巧妙地管理内存是提升处理性能的关键。决定使用哪种格式来处理数据时,你需要了解这些选项是为哪些使用场景设计的。了解这些场景后,你可以在适当的时候利用常见数据传输格式的运行时特性,比如 Protocol Buffers(Protobuf)、JavaScript Object Notation(JSON)和 FlatBuffers。通过合理使用内存,你可以在最小的内存开销下处理海量数据,甚至可以将数据转移到如图形处理单元(GPU)等设备上以获得更高的性能。

在本章中,我们将探讨以下内容:

  • 什么时候使用 Arrow 比使用 Protobuf、JSON、FlatBuffers 或存储格式(如 CSV 和 Apache Parquet)更合适
  • 如何通过内存映射处理巨大的 100 GB 文件,同时只使用少量的物理 RAM
  • 内存映射在幕后做了什么
  • Arrow 库对非 CPU 设备内存的抽象处理

技术要求

和前几章一样,你需要一台联网的电脑,并安装以下软件,以便你可以跟随代码示例一起操作:

  • Python 3.8 或更高版本,已安装 PyArrow 和 pandas(或 Polars)模块
  • Go 1.21 或更高版本
  • 一款支持 C++17 或更高版本编译的 C++ 编译器,且已安装 Arrow 库
  • 你喜欢的 IDE,例如 Emacs、Sublime 或 VS Code
  • 一款网络浏览器
  • Nvidia 显卡及 CUDA 库(可选)

存储格式 vs 运行时内存格式 vs 消息传递格式

当我们谈论用于表示数据的格式时,通常试图优化几个不同但互补的方面。我们可以简化为三个主要组件,如下所示:

  • 大小:最终数据表示的大小。
  • 序列化/反序列化速度:将数据在内存中用于计算的格式之间转换所需的时间性能。
  • 易用性:涵盖可读性、兼容性、功能等方面的一个笼统类别。

我们如何在这些组件之间进行优化,通常取决于该格式的使用场景。在处理数据时,我通常将大多数情况归为三类:长期存储、运行时内存处理和消息传递。虽然这些分类比较广泛,但几乎所有数据的使用都可以归入其中之一。

长期存储格式

什么样的存储格式才是优秀的?通常,数据的大小是首先考虑的因素,因为我们通常需要在单一格式中共享文件和存储大型数据集。随着云存储的兴起与普及,存储成本不再是首要因素,输入/输出(I/O)成本和数据检索的带宽成本以及网络延迟成为了关键因素。无论是使用 Amazon 的简单存储服务(S3)、微软的 Azure Blob 存储,还是类似的服务,虽然存储几乎是无限的且费用极低,但检索大数据时会产生带宽费用和网络延迟。

数据的使用方式不同,在优化读取与优化写入之间通常会存在取舍,因为优化一个通常会对另一个造成不利影响。以下是一些持久存储格式的例子:

  • CSV
  • Apache Parquet
  • Apache Avro
  • Apache Optimized Row Columnar(ORC)
  • JSON

几乎所有情况下,二进制格式都比纯文本格式更小、更高效,因为它可以用更少的字节表示更多的数据。但这通常会牺牲易用性,因为我们无法用眼睛直接读取二进制数据(除非你是机器人)。正因为如此,大多数情况下建议将数据存储为紧凑的二进制格式(如 Parquet),而不是文本格式(如 CSV 或 JSON)。比如,你可以查看本书的代码库中的 sample_data 目录,注意其中包含相同数据的 train.csvtrain.parquet 文件的大小差异。

在查看存储在磁盘上的数据时,有两种方式可以优化读取,如下所示:

  • 减少数据的物理大小
  • 使用一种格式来尽量减少读取满足查询所需的数据量

由于从磁盘读取数据的速度(即便是固态硬盘)远远慢于从内存中读取,因此处理的瓶颈通常在于需要读取多少数据到内存中。虽然这种方式可以使数据从磁盘读取到主内存的效率更高,但在进行内存计算之前仍需解压数据,因此在内存计算场景下并不是最优的选择,如下图所示:

image.png

二进制格式通常更容易压缩,这取决于它们如何存储数据,有时还可以使用一些技巧使数据更具可压缩性。例如,Parquet 和 ORC 采用列式存储形式,通常比 CSV 和 Avro 更具压缩性。你还需要考虑数据的使用方式。虽然列式存储格式可以生成更小的文件,但这只有在常见工作流中只需要行中的一部分列时才最有利。现代大数据中最常见的工作流是在线事务处理(OLTP)和在线分析处理(OLAP)。

OLTP 系统通常以记录(行)级别处理数据,并对每条记录执行创建、读取、更新和删除(CRUD)操作。OLTP 系统执行的任务主要集中在在多用户同时访问时保持数据的完整性,并通过每秒可以处理的事务数量来衡量其有效性。在大多数情况下,OLTP 系统需要的是整个数据记录,而不仅仅是某些字段的子集。对于这种类型的工作流,基于行的存储格式(如 Apache Avro)最为有利,因为通常需要整个记录。

OLAP 系统旨在快速执行分析操作,例如跨多个维度(如时间和其他字段)进行聚合、过滤和统计分析。如前几章所述,这种类型的工作流更适合列式存储,因为这样可以通过只读取所需的数据列来减少 I/O,同时忽略与查询无关的字段。通过将相同类型的数据对齐在一起,你可以获得更高的压缩比,并优化稀疏列中空值的表示。

一旦从存储格式中读取数据,你需要将数据转换为另一种表示形式,以便对其进行操作和计算。这种表示形式称为内存中的运行时表示或格式,如下图所示:

image.png

在许多情况下,I/O 成本远高于将数据解压缩或反序列化为可用的内存表示形式的成本,但这仍然是与数据大小成正比的成本,并且在执行任何操作之前必须支付。当然,也有例外。确实存在一些特定的格式和算法,可以在不解压缩数据的情况下对压缩数据执行操作,但这些通常是非常罕见且专门的情况。

那么,当我们将数据转换为可以操作的某种表示形式时,可以使用哪些格式呢?

内存中的运行时格式

像 Arrow 和 FlatBuffers 这样的格式,其进程间表示与内存表示是相同的,这使得开发人员在将数据从一个进程传递到另一个进程(无论是跨网络还是在同一台机器上的进程之间)时可以避免拷贝数据的成本。属于这一类别的格式的目标是优化计算和运算。

在内存中操作数据时,数据的大小相较于磁盘格式而言不那么重要。在内存中进行计算和分析时,瓶颈主要是中央处理单元(CPU)本身,而不是缓慢的 I/O。在这种情况下,现代开发人员加速性能的主要方法是通过更好的算法使用和优化,例如向量化。因此,能够更好地利用这些优化的数据格式在分析计算中表现得更加高效。这就是为什么 Arrow 设计成当前这种形式的原因,它针对最常见的分析算法进行了优化,并利用了单指令多数据(SIMD)向量化,而不是为磁盘常驻数据进行优化。你可能还记得在第1章《Apache Arrow 入门》中我们讨论过 SIMD。

还有其他因素可能会产生影响,例如字节对齐和随机读取与顺序读取的差异,这些都会对内存处理和磁盘常驻处理的性能产生不同的影响。现代 CPU 并不是一次只执行一条指令,而是使用流水线技术,将指令的执行分为不同阶段,从而错开执行时间。为了最大化吞吐量,CPU 会对即将执行的指令进行各种预测,以保持流水线尽可能地满载运行。下图展示了这种流水线概念的简化示例:

image.png

只要处理器的预测是正确的,执行过程就会顺利进行,所有操作都会以最快的速度完成。下图显示了四条指令在七个时钟周期内完成。当硬件尝试预测 CPU 需要访问的内存位置时,它会在指令执行之前将数据预加载到寄存器中。因此,顺序读取连续的内存块可以非常快速,因为处理器可以轻松预测,并通过较少的指令将更大块的数据加载到主内存中。Arrow 库使用向量化也有助于实现这种高吞吐量,因为这些指令对于现代处理器来说是高度可预测的。当处理器的预测不正确时,就会出现所谓的“气泡”,如下图所示:

image.png

如果执行指令 a 的结果与处理器预测的执行分支不同,流水线将被清空,必须重新从正确的后续指令开始。这会导致与预测正确的情况相比,同样执行4条指令时需要更多的时钟周期。在下图中,可以看到由于错误预测,执行4条指令的时间从7个时钟周期增加到了10个。这虽然是一个简化的例子,但长期来看,大量的错误预测会显著延长运行时间,因为处理器无法最大化其吞吐量。利用这种处理器流水线优化是 Arrow 库的一个重要原则。

硬件对执行的影响还体现在其他方面。与顺序读取相比,传统的旋转磁盘硬盘在随机读取时特别慢,因为它需要移动到另一个磁盘扇区,而不是继续从当前扇区读取。相比之下,固态硬盘 (SSD) 在处理随机读取时几乎没有额外开销。讨论内存表示形式时,操作系统将数据加载到主存中的方式决定了随机读取和顺序读取之间的瓶颈,如我们在第1章《开始使用 Apache Arrow》中所讨论的那样。

内存表示形式中,数据保持在连续的内存块中越多,将这些数据加载到 CPU 的寄存器中就越快。一次性加载大块连续内存的指令比多次加载不同内存区域的指令要快得多。这也是选择列式存储表示形式的原因,Arrow 的这种设计使操作系统和处理器更容易预测执行时所需的内存类型。与此相对,对于某些 OLTP 工作负载来说,像 FlatBuffers 这样的记录导向结构表示可能比列式存储更高效。关键在于理解磁盘格式与临时运行时内存表示之间的权衡和使用场景差异。这就引出了我们讨论的最后一种数据格式——专门为结构化消息传递优化的格式。

消息传递格式

最后要讨论的数据格式类别是消息传递格式,比如 Protobuf、FlatBuffers 和 JSON。在计算机科学中,旨在协调和传递不同进程之间消息的编程接口称为进程间通信 (IPC) 接口。对于消息传递和 IPC,优化小数据包的大小非常重要,因为这些数据通常会通过网络传输。此外,能够轻松地流式传输数据而不丢失上下文信息也是一个巨大的优势。在 Protobuf 和 FlatBuffers 的情况下,开发者在外部文件中定义模式,并使用代码生成器生成处理消息的优化代码。这些技术通常针对较小的消息大小进行优化,处理较大的消息时可能会变得笨重或效率下降。例如,Protobuf 的官方文档指出其最适合的消息大小是小于 1 MB。

IPC 不仅仅是网络

提到 IPC 时,大多数人首先想到的是网络通信以及计算机之间的互联网通信。值得注意的是,IPC 还包括其他一些场景,如同一台计算机上的进程间通信、具有 100 Gbps 互连的云环境以及设备拥有强大专用网络互连卡 (NIC) 的高性能计算 (HPC) 集群中的通信。

虽然这些格式的目的是传递消息,而不是成为为计算优化的格式,但序列化和反序列化的成本仍然至关重要。数据序列化和反序列化的速度越快,数据就能越快发送并进入主存中供计算使用。这也是 Arrow 的 IPC 格式采用原始数据缓冲区而不带有序列化开销的核心原因,但其元数据结构采用 FlatBuffers 来受益于其在消息传递中的高性能。

如果将 Arrow 与 Protobuf 或 FlatBuffers 进行比较,可以看出它们是为解决不同问题而设计的。Protobuf 和 FlatBuffers 提供了一种通用的表示方式,用于通过网络传递消息。通过模式和提供的工具,您可以生成适合所选编程语言的代码,该代码能够接收这种通用的压缩消息格式,并将其转换为可在内存中使用的表示。对于较小的消息,这种序列化和反序列化的开销相对可以忽略不计。而对于需要更高性能的应用场景,FlatBuffers 允许在不需要从反序列化结构中复制数据的情况下操作数据。

然而,当处理非常大的数据集时,序列化和反序列化的成本会大幅增加,成为瓶颈。此外,您无法直接将 C++ 中构建的 Protobuf 结构的原始内存字节交给 Java 程序使用。C++ 需要先将该结构序列化为通用表示,再由 Java 反序列化为其内存表示。这正是 Arrow 格式规范诞生的原因——让不同语言的内存结构保持一致,无需在语言之间进行额外的转换。

总结

如果还不够清楚,简化一下:Arrow 并不是像 Protobuf 这样的消息传递格式的竞争技术,也不是像 ORC 或 Parquet 这样的磁盘格式的竞争者(尽管它也有文件格式)。Arrow 的目标使用场景是补充性的。磁盘格式是为长期、持久存储在磁盘上而设计的,在操作数据之前需要进行大量的解码和解压缩。像 Protobuf 这样的消息传递格式通过使用将整数值打包在一起和使用标签省略默认值的可选字段等技巧来压缩数据。这是为了减少消息的大小,便于在网络上快速传输;如果使用 Arrow 作为消息传递格式,可能会效率较低。

Arrow 设计为一种短暂的、运行时的内存格式,目的是在每个数组单元的基础上处理数据,且开销极小。它并不是为了长期、持久的存储而设计的,尽管为了方便,Arrow 也提供了文件格式。但当你需要传递 Arrow 的记录批和数据表时,你需要了解 Arrow 的 IPC 格式以及如何使用它。

传递你的 Arrows 数据

由于 Arrow 设计为可以在进程之间轻松传递,无论这些进程是否位于同一台机器上,其用于传递记录批(record batches)的接口被称为 Arrow 的 IPC(进程间通信)库。如果这些进程位于同一台机器上,甚至可以在不进行任何数据复制的情况下共享数据!

这是什么神奇的操作?!

首先,这里定义了两种用于在进程间共享记录批的二进制格式:一种是流式格式,另一种是随机访问格式,详细介绍如下:

  • 流式格式:用于发送任意长度的记录批序列。它必须从头到尾依次处理,不能随机访问流中的某个记录批,而必须处理其之前的所有记录批。
  • 随机访问格式(也称为文件格式):用于共享已知数量的记录批。它支持对特定记录批的随机访问,非常适合与内存映射结合使用。

记住
要理解发送二进制 Arrow 数据的格式,关键是要记住什么是记录批。记录批只是一些有序的数组集合,这些数组长度相同但数据类型可能不同。这些字段的名称、类型以及元数据统称为记录批的模式(schema)。

Arrow 的 IPC 协议定义了三种用于传递信息的消息类型:SchemaRecordBatchDictionaryBatch。这些消息的二进制负载序列可以在无需内存复制的情况下重构为内存中的记录批。每条消息由一个 FlatBuffers 格式的元数据消息和一个可选的消息体组成。FlatBuffers 是一个由 Google 原创设计的高效跨平台序列化库。FlatBuffers 的设计允许消息直接解释和访问,而无需先反序列化为另一种中间格式。(你可以在 FlatBuffers 官方网站 学习更多相关内容。)

下图展示了沿流发送的封装消息格式。前 8 字节由 0xFFFFFFFF 组成,表示这是一个有效消息,后跟 4 字节的小端序(LE)整数值,表示 FlatBuffers 消息和填充的大小。最后,消息体由一个可选的原始字节序列组成,其长度必须是 8 字节的倍数。

image.png

好的,你可能会想:“嘿,我不是刚刚读到这里没有反序列化的开销吗?难道我不需要解包 FlatBuffers 数据吗?”实际上,选择 FlatBuffers 的原因是可以直接访问 FlatBuffers 消息中的二进制数据,而无需将其解包为其他格式。只需要访问消息本身的内存,不需要额外的内存分配。不仅如此,FlatBuffers 还允许只将部分缓冲区加载到内存中,而无需一次性加载全部内容。总结一句话?它非常快! 另外请记住,只有当你在实现这个规范时才需要处理 FlatBuffers 部分;如果不是,你可以放心交给 Arrow 库处理这一切。

接下来我们来看一下 FlatBuffers 消息表的结构:

table Message {
    version: org.apache.arrow.flatbuf.MetadataVersion;
    header: MessageHeader;
    bodyLength: long;
    custom_metadata: [ KeyValue ];
}

如你所见,FlatBuffers 消息数据包含一个格式版本号、一个特定的消息值(可以是 Schema、RecordBatch 或 DictionaryBatch)、以字节为单位的消息体长度,以及一个用于应用定义的键值对元数据的字段。一般来说,读取消息流的过程是首先读取 FlatBuffers 消息值以获取消息体的大小,然后读取消息体的字节数据。典型的消息流通常由一个 Schema 消息开始,后跟一些 DictionaryBatchRecordBatch 消息。首个 Schema 消息不包含任何数据缓冲区,仅包含关于消息的元数据信息,例如类型信息。我们可以将数据流可视化为如下所示:

image.png

每个消息流中的框代表一个封装的消息,包含了 FlatBuffers 数据和可选的消息体缓冲区。如前所述,字典消息只有在模式中包含字典编码的数组时才会出现。每个字典消息都有一个 id 字段,该字段被模式引用,用来指示哪些数组使用该字典。如果可能,通过引用多个数组的 id 字段,可以优化使用相同字典的多个字段。此外,还有一个 isDelta 标志,允许在后续的记录批次中扩展现有字典,而无需重新发送整个字典。Schema 消息不会包含任何消息体数据,它们只包含描述类型和元数据的 FlatBuffers 消息。

快速提示
这个简短部分的其余内容介绍了 IPC(进程间通信)实现的一些底层技术细节。虽然这对于理解其工作原理很重要,但并非必须掌握这些细节才能利用 Arrow IPC 格式及其相关功能。

Record-batch 消息包含以下内容:

  • 一个由 RecordBatch FlatBuffers 消息定义的头部,包含每个字段的长度和空值计数,以及消息体中每个对应数据缓冲区的内存偏移和长度。为了处理嵌套类型,字段会被扁平化成按深度优先遍历的顺序。

  • 例如,考虑以下模式:

    Col1: Struct<a: int32, b: List<item: float32>, c: float64>
    Col2: String
    

扁平化版本会像这样:

  • Field0: Struct name='Col1'
  • Field1: Int32 name='a'
  • Field2: List name='b'
  • Field3: Float32 name='item'
  • Field4: Float64 name='c'
  • Field5: UTF8 name='Col2'

这些记录批次由从头到尾的原始数据缓冲区组成,并且进行了填充以确保 8 字节对齐:

  • 对于记录批次中的每个扁平化字段,组成数组的数据缓冲区将根据我们在第1章《初识 Apache Arrow》中介绍的描述来决定:包括有效性位图、原始数据、偏移量等等。

在读取消息时,无需复制或转换数据缓冲区;它们可以直接被引用,不需要进行任何复制或反序列化。

在每条消息之后,读取器可以读取接下来的 8 个字节以确定是否有更多数据,并获取下一个 FlatBuffers 元数据消息的大小。写入器可以通过以下两种方式之一来表示没有更多数据(即结束流,EOS):

  1. 关闭流
  2. 写入 8 个字节,这些字节包含一个 4 字节的继续标识符(0xFFFFFFFF),然后是表示下一个消息长度的 4 字节整数(0x00000000)

在我们展示如何使用 IPC 流的示例之前,还有一件事要介绍,那就是随机访问文件格式。相信我,如果你理解了流格式,文件格式将非常容易理解。下图描述了文件格式:

image.png

文件格式只是流格式的扩展,文件的开头和结尾都有一个魔法字符串指示符,并且包含一个页脚。页脚中包含模式的副本(记住,流格式的第一个消息也是模式)以及文件中每个数据块的内存偏移量和长度,从而支持文件中任意记录批次的随机访问。建议使用 .arrow 文件扩展名。流通常不会写入文件,但如果写入,建议使用 .arrows 扩展名。此外,Apache Arrow 数据的流格式和文件格式都注册了多用途互联网邮件扩展(MIME)媒体类型,如下所示:

额外信息 IPC 格式中有一个选项表明可以对各个主体缓冲区进行压缩,以进一步减少消息传输时的大小,但代价是增加了一些 CPU 使用。这在网络延迟是瓶颈而不是 CPU 的情况下非常有用,或者在写入文件时为了减少文件大小也很有帮助。IPC 格式支持的两种压缩类型是 Zstandard(ZSTD)和 Lempel-Ziv 4(LZ4)压缩。

好了,我们已经介绍了传输 Arrow 数据的协议,现在来看一些示例。

生成和消费 Arrow 数据

如果你想构建一个服务来读取远程文件,并将数据流式传输回客户端,应该怎么做?听起来很简单,但如果在读取数据文件和传输数据之间进行一些数据处理,这实际上就是很多数据传输用例的基础。因此,我们将构建类似的服务:

image.png

通过遵循图 3.8 的步骤,以下是具体操作步骤:

  1. 消费方请求一个文件(这里不详细说明,但会在 GitHub 示例中提供)。
  2. 从 S3 建立输入流,将文件读取为 Arrow 记录批次。
  3. 在从 S3 读取记录批次时,将其写入给消费方。

我们来试试如何实现这个过程吧!我们之前已经介绍了如何在 Python 和 C++ 中从 S3 读取文件,因此我们来用 Go 语言编写生产者端的示例。

首先,我们需要导入所需的包,如下所示:

import (
    "compress/gzip"
    "context"
    "github.com/apache/arrow/go/v17/arrow"
    "github.com/apache/arrow/go/v17/arrow/csv"
    "github.com/apache/arrow/go/v17/arrow/ipc"
    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/service/s3"
)

接下来,我们可以创建 S3 客户端并建立输入流,如下所示:

client := s3.New(s3.Options{Region: "us-east-1"})
obj, err := client.GetObject(context.Background(),
      &s3.GetObjectInput{
          Bucket: aws.String("dataforgood-fb-data"),
          Key: aws.String("csv/month=2019-06/country=ZWE/type=children_under_five/ZWE_children_under_five.csv.gz")})
if err != nil {
   // 处理错误
}

假设没有错误,返回的 obj 对象的 Body 成员将是一个流,用于读取文件的数据。不过,在这个示例中,返回的文件是用 gzip 压缩的。因此,我们需要创建一个解压缩的读取器,如下所示:

rdr, err := gzip.NewReader(obj.Body)
if err != nil {
    // 处理错误
}
defer rdr.Close()

现在,我们可以设置 CSV 读取器,如下所示:

schema := arrow.NewSchema([]arrow.Field{
    // 预期的 CSV 文件的架构
}, nil)
headerline := true // 如果第一行不是列头,设为 false
delim := '\t' // 文件使用制表符分隔而非逗号
reader := csv.NewReader(rdr, schema,
    csv.WithComma(delim), csv.WithHeader(headerline))
defer reader.Release()

你可以阅读文档,了解创建 CSV 读取器时的其他可能选项。

为了设置 IPC 流,我们创建一个 ipc.Writer 接口并传递一个流来写入。这个流可以是任何实现了 io.Writer 接口的对象,无论是文件还是 HTTP 响应写入器。代码如下所示:

writer := ipc.NewWriter(outStream, ipc.WithSchema(
              reader.Schema()))
defer writer.Close()

最后,我们可以一边读取记录批次,一边将其写出!流程如下所示:

for reader.Next() {
    if err := writer.Write(reader.Record()); err != nil {
                   // 处理错误
       break
    }
}
if reader.Err() != nil {
    // 如果读取器遇到错误停止,而不是因为完成
    // 处理错误!
}

就这样!

消费这个流在任何编程语言中也都非常简单。在读取和写入 IPC 格式时,概念上都是一样的:

  • 在 Go 中,你可以使用 ipc.Reader
  • 在 Python 中,可以使用 pyarrow.RecordBatchStreamReader 进行读取,用 pyarrow.RecordBatchStreamWriter 进行写入。
  • 如果你想读取多个记录批次并将其转换为单个 DataFrame,可以使用 read_pandas 函数简化此过程。

或者,如果你使用的是 Polars,有几个选项可供选择:

  • read_ipc_stream:将 IPC 流读入 DataFrame。
  • scan_ipc:延迟读取一个或多个 IPC 文件到 DataFrame 中,只有在请求结果时才会真正执行读取。
  • DataFrame.write_ipc_stream:将 DataFrame 写入 IPC 流。

C++ 库提供 arrow::ipc::RecordBatchStreamReader::Open,它接受 arrow::io::InputStream 来创建读取器,并提供 arrow::ipc::MakeStreamWriter 来从 arrow::io::OutputStream 创建写入器。

如果你更喜欢使用随机访问文件格式,那么上面的每个函数都有对应的版本用于读取和写入文件格式。如果你曾遇到过函数或文档提到 Feather 文件格式,它只是 Arrow IPC 文件格式在磁盘上的表示,最早是作为语言无关的 DataFrame 存储格式的概念验证而创建的。

为了在处理这些流和文件时确保高效的内存使用,Arrow 库提供了各种用于内存管理的实用工具,它们也在内部使用。我们将讨论这些辅助类和实用工具的使用,以及如何在不同编程语言之间共享缓冲区以提高性能。

Arrow IPC 文件格式通过内存映射提供性能提升。由于 Arrow IPC 格式具有与内存中的数据缓冲区相同的格式,可以利用一种称为内存映射的技术来最小化处理时的内存开销。我们接下来将介绍这种技术及其工作原理。

学习内存映射

分布式系统如 Apache Spark 的吸引力之一在于其能够快速处理超大规模数据集。有时,数据集太大,甚至无法全部加载到单台机器的内存中!在这种情况下,需要分布式系统将数据拆分成块并并行处理,因为没有任何单个机器能够一次性加载整个数据集进行操作。但如果可以在使用几乎不占用 RAM 的情况下处理数 GB 的大文件呢?这就是内存映射的用武之地。

让我们再次使用纽约出租车数据集来帮助演示这个概念。我们下载 2015 年 1 月的黄牌出租车数据,文件名为 yellow_tripdata_2015-01.parquet,大约 168 MB。如果将其转换为 CSV 文件,大小约为 1.3 GB,非常适合作为示例。为了简洁起见,接下来的示例我们将使用 Python Arrow 库 PyArrow。我们使用 CSV 文件作为基线,因为它是最常见的数据文件格式之一。现在,假设我们想要查找并计算 total_amount 列中的平均值。

获取数据

你可以从纽约市政府的公共网站下载我们使用的示例文件:纽约出租车数据。除了 PyArrow,还确保你安装了 IPython,因为我们将使用它来衡量示例的性能。可以通过 pip install ipython 轻松安装。

基线案例

你下载了 yellow_tripdata_2015-01.parquet 文件了吗?如果没有,请赶紧下载!如果愿意,可以试着用我们之前介绍的 PyArrow 库和文档将其转换为 CSV 文件。事不宜迟,启动 IPython 解释器:

In [1]: import pyarrow as pa
In [2]: import pyarrow.csv
In [3]: import pyarrow.parquet as pq
In [4]: table = pq.read_table('yellow_tripdata_2015-01.parquet')
In [5]: pa.csv.write_csv(tbl, 'yellow_tripdata_2015-01.csv')

为了看看内存映射带来的好处,首先我们不使用内存映射计算 total_amount 列的平均值,以获得运行时和内存使用情况的基线。如果我们天真地使用 pandas 来读取 CSV 文件,并告诉它只关心这一列,再计算平均值,这很简单:

In [6]: import pandas as pd
In [7]: %timeit pd.read_csv('yellow_tripdata_2015-01.csv', usecols=['total_amount'])['total_amount'].mean()
3.47 s ± 44.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

由于我们使用了特殊的 %timeit 前缀,IPython 会多次运行该行代码,并平均计算运行所需的时间。我们看到,使用 pandas 的这个简单方法大约花费了 3.47 秒,误差约为 44 毫秒。我们能做得更快吗?如果直接使用 PyArrow 呢?

In [8]: import pyarrow.compute as pc
In [9]: %%timeit
   ...: pc.mean(
   ...:   pa.csv.read_csv('yellow_tripdata_2015-01.csv',
   ...:   convert_options=pa.csv.ConvertOptions(
   ...:      include_columns=['total_amount'])
   ...: )['total_amount'])
   ...:
325 ms ± 10.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

你可以看到 PyArrow 包含自己的计算模块。尽管我们会在后续章节中更详细地介绍它,但现在我们只用它来快速计算平均值,而不必先转换为 pandas DataFrame。正如你所看到的,速度提升巨大!从大约 3 秒半减少到 325 毫秒,快了大约 10 倍!

Parquet 与 CSV 对比

正如我之前提到的,二进制格式可以提供更高的压缩率。在这种情况下,1.3 GB 的 CSV 文件在 Parquet 格式下仅为 168 MB。让我们看看 Parquet 格式的较小文件大小和 I/O 优势能否帮助我们提高计算性能。首先,我们用 pandas 进行简单的计算:

In [10]: %timeit pd.read_parquet('yellow_tripdata_2015-01.parquet', columns=['total_amount'])['total_amount'].mean()
210 ms ± 10.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

然后,我们使用 PyArrow 的计算模块而不使用 pandas:

In [11]: %timeit pc.mean(pq.read_table('yellow_tripdata_2015-01.parquet', columns=['total_amount'])['total_amount'])
178 ms ± 2.15 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

可以理解,使用 pandas 和 PyArrow 在这里的性能大致相当,因为 pandas 也使用 PyArrow 的 Parquet 读取器。

尽管在使用 CSV 文件的例子中我们用列筛选,但它仍然需要解析所有列以找到我们想要的列。而 Parquet 格式则允许我们只读取所需的列。结合压缩数据的较小体积以及仅读取单列数据所减少的 I/O,结果表明,使用 pandas 时需要大约 210 毫秒,而直接使用 PyArrow 则更快,仅为 178 毫秒。

image.png

接下来我们尝试使用内存映射文件。

将数据映射到内存中

稍后我会详细说明内存映射背后的机制,但现在我们先看看如何使用 Arrow 库实现内存映射。为了从内存映射中获益,我们需要将文件写为 Arrow IPC 文件。这是因为,即使我们对 CSV 或 Parquet 文件进行内存映射,仍然需要解码或解析这些格式,因此仍然需要将文件读取到主内存中。我们可以使用 RecordBatchFileWriter 将表格写入文件,代码如下:

In [12]: table = pq.read_table('yellow_tripdata_2015-01.csv').combine_chunks() # 我们想要一个连续的表格
In [13]: with pa.OSFile('yellow_tripdata_2015-01.arrow', 'wb') as sink:
   ...:      with pa.RecordBatchFileWriter(sink, table.schema) as writer:
   ...:         writer.write_table(table)
   ...:

注意我们使用了 .arrow 扩展名来表示文件。这符合由互联网分配号码机构(IANA)认可的标准媒体类型。默认情况下,该文件没有进行任何压缩,因此文件大小相对较大。虽然 CSV 文件大约为 1.3 GB,但 Arrow IPC 文件的大小大约为 1.6 GB,比 CSV 文件多了 300 MB 左右。然后,我们可以将文件映射到内存中,并重复之前的操作来从 Arrow 表格中计算 total_amount 列的平均值,代码如下:

In [14]: import pyarrow.ipc
In [15]: %%timeit
    ...: src = pa.memory_map('yellow_tripdata_2015-01.arrow')
    ...: col = pa.ipc.RecordBatchFileReader(src).read_all().column('total_amount')
    ...: pc.mean(col)
    ...:
7.66 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

哇!这是一个非常夸张的速度提升,对吧?我们从大约 180-200 毫秒降低到了仅 7.66 毫秒。让我们用这个新数据点更新我们的表格。

image.png

为了进一步观察内存映射的好处,我们需要检查每种方法分配和使用了多少内存。为此,我们需要安装 psutil 包,可以通过以下命令安装:

$ pip install psutil

要检查这些方法的内存使用情况,我们可以在每次读取数据之前和之后检查进程的当前内存分配量,并记录差异。我们可以编写一个脚本来完成这个任务,而不是使用 IPython 进行交互式操作。让我们逐步讲解这个脚本,从第一行开始:

# 获取初始内存使用量(单位:MB)
memory_init = psutil.Process(os.getpid()).memory_info().rss >> 20

我们使用 psutil 库检查进程的当前内存使用量,并将值从字节转换为 MB,方便读取。接下来,我们可以读取数据,并通过不同的方式获取我们想要的列,记录每次操作后的内存使用情况。首先,我们通过读取 CSV 文件并提取我们需要的列来做测试,如下所示:

col_pd_csv = pd.read_csv('yellow_tripdata_2015-01.csv', usecols=['total_amount'])
col_pd_csv['total_amount'].mean()
memory_pd_csv = psutil.Process(os.getpid()).memory_info().rss >> 20

接下来,我们使用 PyArrow 的 CSV 读取器:

col_pa_csv = pa.csv.read_csv('yellow_tripdata_2015-01.csv', 
                             convert_options=pa.csv.ConvertOptions(include_columns=['total_amount']))
pc.mean(col_pa_csv['total_amount'])
memory_pa_csv = psutil.Process(os.getpid()).memory_info().rss >> 20

然后,我们使用 Parquet 文件来代替 CSV 文件:

col_parquet = pq.read_table('yellow_tripdata_2015-01.parquet', columns=['total_amount'])
pc.mean(col_parquet['total_amount'])
memory_parquet = psutil.Process(os.getpid()).memory_info().rss >> 20

接着,我们使用 Arrow IPC 文件,但不使用内存映射:

with pa.OSFile('yellow_tripdata_2015-01.arrow', 'rb') as source:
    col_arrow_file = pa.ipc.open_file(source).read_all().column('total_amount')
pc.mean(col_arrow_file)
memory_arrow = psutil.Process(os.getpid()).memory_info().rss >> 20

值得注意的是,使用 read_all 需要将整个 .arrow 文件加载到内存中,才能对列进行过滤。而通过内存映射 Arrow IPC 文件,可以直接引用文件中包含列数据的位置,而不需要加载整个文件。使用内存映射读取列的代码如下:

source = pa.memory_map('yellow_tripdata_2015-01.arrow', 'rb')
col_arrow_mmap = pa.ipc.RecordBatchFileReader(source).read_all().column('total_amount')
pc.mean(col_arrow_mmap)
memory_mmapped = psutil.Process(os.getpid()).memory_info().rss >> 20

代码中的每一行分别存储了每次读取数据时的内存使用情况。我们可以通过逐步相减这些值来查看进程为每种方式分配的内存量。脚本中,各种读取数据方式的内存使用量分别是:

  • 使用 pandas 读取 CSV 文件:memory_pd_csv - memory_init
  • 使用 PyArrow 读取 CSV 文件:memory_pa_csv - memory_pd_csv
  • 读取 Parquet 列:memory_parquet - memory_pa_csv
  • 常规读取 Arrow 文件:memory_arrow - memory_parquet
  • 使用内存映射读取 Arrow IPC 文件:memory_mmapped - memory_arrow

真相揭晓! 各种方法使用的内存是多少呢?让我们看看结果。

image.png

在我们分析结果之前,有一个事实需要牢记:这个数据集包含 12,741,035 行数据,total_amount 列的数据类型是 float64(64 位浮点数值,即每个值占用 8 字节)。如果我们计算一下,表示这列数据所需的内存大约是 97 MB。虽然有一些方法可以在不将整个列加载到内存中的情况下计算平均值,但我们在这里并没有使用这些方法(为了简单起见)。

现在我们得出了结论——即使这个内存映射文件大小为 1,617 MB,但我们实际上只分配了正好能容纳我们所需列的数据的内存。简直像魔法一样!我们还可以看到,通常读取 Arrow IPC 文件时的内存使用量略小于整个文件大小,因为它需要将整个文件加载到内存中进行处理。为什么会略小呢?这很可能是因为 IPC 文件中的页脚和其他 FlatBuffers 消息并不需要复制即可处理文件。与此同时,Parquet 文件仅使用了非常合理的 123 MB 内存。而真正令人震惊的是,使用 PyArrow 库读取 CSV 文件时,内存使用量甚至远远超过了 CSV 文件本身的大小。这是因为多线程读取不会分配一个连续的大内存块,而是同时分配多个较小的内存块,提前预分配缓冲区。

内存映射文件方法比 Parquet 版本更快有两个原因:只需分配我们需要的列所需的内存量,并且不需要进行解码来访问数据。我们只需直接引用内存并可以立即操作它!

公平地说…

在这个场景中,解码成本很可能比内存分配成本占据更重要的时间影响。这是因为解码和复制的成本通常比这里发生的“小规模”内存分配要高得多。

如果你对这种神奇的魔法感到好奇,那就继续阅读吧,我亲爱的冒险读者!

长话短说(TL;DR)——计算机就是魔法

在我解释内存映射背后的“魔法”之前,你需要了解运行进程的内存是如何工作的。这里会有些简化的内容,这没问题!这是一次速成课,而不是大学课程。

虚拟和物理内存空间

大多数现代操作系统使用虚拟化内存的概念来管理进程使用的内存。这意味着进程可能引用的地址空间并不需要全都存储在设备的物理 RAM 中。相反,操作系统可以将进程的内存数据从主内存中交换出去,写入页面文件或缓存,或者以其他它希望的方式处理。在这种上下文中,可以将一个数据页面视为一块内存,通常大小约为 4 KB。操作系统维护一个映射表,称为页面表,用于将虚拟地址映射到存储该数据的物理 RAM 位置。如果请求加载的数据不在物理 RAM 中,则称之为页面错误。然后,所请求的数据会被交换回 RAM 中,以便进程可以正确引用。这一切对进程来说是透明的,完全由操作系统处理。下图展示了这个过程的示意图:

image.png

图 3.12 展示了一个进程的虚拟内存视图如何映射到物理内存和磁盘的示意图。进程可以将其内存视为一个连续的空间,其中包含其栈、堆、共享库及其内存中需要引用的任何其他内容。但实际上,这些内存块可能位于 RAM 芯片的不同位置,甚至可能在页面文件或某个地方的缓存中。操作系统负责跟踪虚拟内存位置和物理内存位置之间的映射,并在进程需要时将数据加载到物理 RAM 中或从中移出时更新该表。

虚拟化内存有许多好处,使其成为事实上的标准,以下是其中一些:

  • 进程不需要自己管理共享内存空间
  • 操作系统可以在进程之间共享相同的内存空间,从而利用共享库并减少整体内存消耗
  • 通过将一个进程的内存与另一个进程的内存隔离,增强了安全性
  • 可以在概念上引用和操作比实际可用的物理内存更多的内存

当从文件中读取数据时,进程会进行系统调用,每次读取或写入文件时都需要经过不同的驱动程序和内存管理器。在许多步骤中,数据会被缓冲和缓存(即被复制)。而内存映射文件允许我们绕过所有这些拷贝操作,并将文件的数据视为普通内存进行加载/存储操作,而不是使用系统调用,后者的速度要慢得多。

什么是内存映射?

内存映射是由 POSIX 定义的 mmap 函数提供的一项功能,在许多操作系统上都有。当你对一个文件进行内存映射时,进程的虚拟地址空间中与文件大小等同的一块区域被分配出来,从而允许进程将文件数据视为简单的内存加载和存储操作。在大多数情况下,映射的内存区域是内核的页面(文件)缓存,这意味着访问数据时不会创建额外的副本。

下图展示了文件被内存映射时的一个示例。当所有对文件的访问都通过内核的页面缓存时,文件的更新会立即对试图读取文件的任何进程可见,尽管这些更新在缓冲区刷新并同步回物理驱动器之前不会物理存在于设备上。这意味着多个进程可以对同一个文件进行内存映射,并将其视为共享内存区域,这也是底层实现专用共享进程内存的方式。

image.png

在我们将大型 Arrow 文件内存映射以读取值并处理列的示例中,我们利用了内存映射和 PyArrow 库的延迟加载和零拷贝特性。通过对 1.6 GB 文件进行内存映射,该库能够将文件视为已经在内存中,而无需实际分配整个 1.6 GB 的文件空间。它只会在访问相应的虚拟内存位置时,从文件中读取所需的页面数据。与此同时,我们可以传递指向内存地址的引用。我们从文件底部读取了几 KB 的数据,正如你还记得的那样(在第 1 章《初识 Apache Arrow》中提到),那里存储了元数据和架构,只访问相应的页面。当我们最终想要读取列中的值以进行计算时,数据才会被实际加载并拉入 RAM。回头看看“将数据映射到内存”部分的代码片段。

如果我们删除执行均值计算的高亮代码行,那么你会看到内存使用量完全消失了。由于我们从未访问到列数据的内存部分,它只会加载几 KB 的元数据页:

source = pa.memory_map('yellow_tripdata_2015-01.arrow', 'rb')
col_arrow_mmap = pa.ipc.RecordBatchFileReader(source).read_all().column('total_amount')
pc.mean(col_arrow_mmap)
memory_mmapped = psutil.Process(os.getpid()).memory_info().rss >> 20

带有该行的脚本运行时,会显示我们之前看到的 97 MB——即整个列的大小——但如果我们删除该行并再次运行,它现在报告为 0 MB 的使用量。相当酷吧?

那么,为什么我们不总是对文件进行内存映射呢?其实并不是这么简单的。让我们看看原因。

内存映射并非万能

正如我们在前面几节中看到的,使用持久性内存映射文件的主要原因是为了提高 I/O 性能。但和其他技术一样,这也是一种权衡。虽然标准的 I/O 方法由于系统调用和内存拷贝的开销而比较慢,但内存映射的 I/O 也有其成本:次级页面错误。当进程试图访问尚未加载到内存的虚拟内存空间页面时,就会发生页面错误。对于内存映射 I/O,当页面已存在于内存中,但系统的内存管理单元尚未标记其为已加载,而进程访问时发生的就是次级页面错误。这种情况发生在数据块已经加载到页面缓存中(如图 3.13 所示),但尚未映射并连接到进程虚拟内存空间中的适当位置。在某些情况下,取决于访问模式和硬件,内存映射 I/O 可能比标准 I/O 慢得多。

对于非常小的文件(大约为 KB 或更小),对它们进行内存映射可能会导致内存碎片过多,浪费内存空间。内存映射区域总是对齐到内存页面大小,而在大多数系统中,页面大小通常为 4 KB。一旦整个文件加载到内存中,映射一个 5 KB 的文件将需要两个内存页面,即 8 KB 的分配内存。下图展示了如何导致内存中 3 KB 的浪费闲置空间:

image.png

虽然 3 KB 看起来并不是很多内存空间,但如果你对大量小文件进行内存映射,这些浪费和碎片化的内存会迅速积累。最后,如果你使用的是非常快的设备,比如现代的 NVMe SSD,可能会受到操作系统限制处理页面错误的 CPU 核心数量的影响。这导致内存映射 I/O 比标准 I/O 可扩展性更差。如果有太多进程执行内存映射操作,瓶颈就会出现在核心数量的限制上,而不是存储设备的性能上。

使用内存映射文件时,还需要考虑以下几点:

  • 如果在访问映射内存时,支持源文件发生任何 I/O 错误,进程会将其报告为错误。这些错误在普通内存访问中通常不会发生。因此,访问映射内存的代码需要准备好处理这些错误。
  • 在没有内存管理单元的硬件上(如嵌入式系统或低功耗系统),当请求映射文件时,操作系统可能会将整个文件复制到内存中。这显然只适用于可以装入可用内存的文件,而如果你只需要访问文件的一小部分,这样的操作会非常慢。
  • 正如前面提到的,使用内存映射文件时,大部分操作都发生在触发主要页面错误时,比如尝试访问内存时。这些错误可能会在运行时的任何时刻发生(如解引用指针时),因此无法预见或调度它们的发生时间。结果是,I/O 成本和计算成本无法分离,导致扩展性差,并且难以利用内存映射进行异步 I/O。

总之,内存映射是处理非常大的文件(尤其是只需读取其中一小部分)以及在进程间共享内存的重要技术。只要确保在使用时频繁进行测试和基准测试,以确定它是否真正带来了效益。需要记住的是,内存映射诸如 Parquet 文件时,效益会比映射原始 Arrow IPC 文件低。你可能会问:为什么?因为在你使用 Parquet 文件中的数据之前,仍然需要对其进行解码和解压缩,内存仍然需要分配以存储解码后的数据。虽然你可能在 Parquet 文件的 I/O 成本上节省了一点性能,但对于速度足够快的磁盘,瓶颈会出现在内存分配和解码/解压缩所需的 CPU 上。相比之下,Arrow IPC 文件的原始字节可以直接在内存中使用,不需要额外的内存分配来引用和使用文件中的数据,这就是它的优势。

到目前为止,我们讨论的所有内存相关内容都只涉及 CPU 可访问的内存,即主内存。然而,随着越来越多的设备(如 GPU)因其在矢量化和分析计算方面的显著性能优势而被广泛用于计算,Arrow 也为此做好了准备!

离开 CPU——使用设备内存

随着大语言模型(LLM)和其他机器学习(ML)用例的普及,越来越多的库和工作流开始利用 GPU 或其他硬件设备。这种转变要求我们采用全新的工程范式,这可能很难学习和应用。为帮助这一过渡,Arrow C++ 库——以及 PyArrow 作为其扩展——提供了一系列接口和构件,用于设计能够在主内存和设备内存中使用 Arrow 格式数据的系统。

重要提示!

在本节和后续章节中,我们会有一些关于设备内存(特别是 Nvidia 图形卡和 CUDA 架构)示例,但我不会深入讨论这些设备的具体编程。很多在这方面经验丰富的人已经编写了整本书来教授这些内容。本书的重点是 Arrow 格式数据与流行库之间的互操作性。

从一些指针开始

明白了吗?"Arrows" 代表箭头,而指针也称为 memory pointers……?继续说正事。 到目前为止,我们的所有代码示例都使用了默认的内存管理,大多数情况下是隐式的。但是,我所涉及的三种语言(C++、Python 和 Go)都提供了自定义内存分配和管理的方式。通常,这些被称为“内存池”(C++ 和 Python)或“分配器接口”(Go)。默认情况下,这些内存池和分配器是为 CPU 可访问的主内存服务的。

为了表示不同的设备,Arrow 提供了一个名为 Device 的抽象对象。每个 Device 都有一个默认的 MemoryManager 对象,用于指定如何在该设备上分配内存。同一设备上可能还有额外的内存管理器对象,这些对象可能会封装自定义的内存池。

这些抽象对象被用于 Arrow 的缓冲区代码中,使得一个给定的 Buffer 可以持有与任意设备相关的内存。如果你从第三方代码接收到 Arrow 数据,且这些数据可能是在非 CPU 设备上分配的,那么你应该检查 Buffer 对象的 is_cpu() 方法。如果返回 true,则数据是直接 CPU 可访问的。

为什么使用 Arrow 与 GPU 协作?

GPU 设备非常擅长进行矢量化计算操作。通常,这类似于,或者直接使用了我们之前讨论过的 SIMD 处理(单指令多数据流)。虽然有些 GPU 架构不直接使用 SIMD,但它们仍然针对类似的矢量操作进行了优化。

设备无关的缓冲区处理

上面的内容有些复杂?别担心,没关系! Buffer 对象有一些方法可以帮助你以设备无关的方式编写代码:

  • arrow::Buffer::View:给定一个缓冲区和目标内存管理器,设备特定机制会生成一个内存地址,用于在目标设备上访问缓冲区的内容。如果无法在不复制的情况下实现这一点,则会返回错误状态。如果源设备和目标设备相同,则不会进行任何操作。
  • arrow::Buffer::Copy:将缓冲区内容复制到目标设备,并由提供的内存管理器分配一个新缓冲区。
  • arrow::Buffer::CopyNonOwned:这是 Copy 的替代方法,适用于源缓冲区由外部管理且其生命周期不受 Buffer 对象控制的情况。否则,应该使用 Copy
  • arrow::Buffer::ViewOrCopy:首先尝试在目标设备上创建一个零拷贝视图,如果无法做到,则将内容复制到一个新缓冲区。

对于上述方法,设备之间的实际数据传输可能会延迟,直到访问缓冲区内容时才发生。

下面是一个确保缓冲区在 CPU 上可访问的代码示例:

std::shared_ptr<arrow::Buffer> arbitrary_buf = ...;
std::shared_ptr<arrow::Buffer> cpu_buf;
auto maybe_cpu_buf = arrow::Buffer::ViewOrCopy(arbitrary_buf, arrow::default_cpu_memory_manager());
if (!maybe_cpu_buf.ok()) {
    ARROW_LOG(ERROR) << maybe_cpu_buf.status();
} else {
    cpu_buf = maybe_cpu_buf.MoveValueUnsafe();
}

另一个常见操作是将缓冲区传递给 I/O 函数,这些函数可能直接从缓冲区读取或写入。如果你想进行 I/O 操作,但不假设缓冲区是 CPU 可访问的,可以使用 arrow::Buffer::GetReaderarrow::Buffer::GetWriter 函数来获取设备特定的流接口。

以下是如何使用 Arrow 接口向设备写入数据的示例:

std::shared_ptr<arrow::MemoryManager> mm = ...;
auto buf = mm->AllocateBuffer(10).ValueOrDie();
auto writer = arrow::Buffer::GetWriter(buf).ValueOrDie();
auto status = writer->Write(reinterpret_cast<uint8_t*>(
                                "some data!"), 10);
if (!status.ok()) {
    ARROW_LOG(ERROR) << status;
}

之后,我们可以使用 arrow::Buffer::GetReader 来读取数据,不论底层设备如何:

auto reader = arrow::Buffer::GetReader(buf).ValueOrDie();
if (reader->GetSize() != 10) {
    ARROW_LOG(ERROR) << "wtf??";
}
auto read_buf = reader->ReadAt(5, 5).ValueOrDie();
if (read_buf->ToString() != "data!") {
    ARROW_LOG(ERROR) << "wtf??";
}

这些接口可能听起来有点抽象,但它们非常实用。在接下来的章节中,我们会通过一些实际的例子展示如何在 ML 工作流中利用这些接口。

总结

无论你是数据科学家还是软件架构师,在构建数据管道和大型系统时,你都需要做出许多关于选择何种格式的决策。你总是希望根据具体的用例选择最合适的格式,而不是随意追随潮流并将其应用到所有地方。许多人在听到 Arrow 时,或者反应过度,认为他们需要在所有地方都使用它,或者疑惑为什么我们需要另一个数据格式。关键的一点是要理解这些格式所试图解决的问题的不同之处。

如果你需要长时间的持久存储,无论是在本地磁盘还是在云端,通常你会选择像 Parquet、ORC 或 CSV 这样的存储格式。在这些用例中,主要的访问成本是 I/O 时间,因此你需要根据访问模式来优化以减少这种成本。如果你只是传递小消息,比如元数据或控制消息,那么 Protobuf、FlatBuffers 和 JSON 这样的格式可能是最优的选择。然而,这些格式并不适用于大规模的表格数据集,尤其是当你需要对数据进行分析和计算时。这里没有硬性规定——这些只是指导原则。Arrow 针对的用例是作为一种内存中短暂的运行时格式,为数据提供共享内存格式,使处理时间更多地用于执行计算,而不是在数据在系统中传递时花费大量时间在格式转换上。

希望你也学到了在处理这些大型文件和数据时如何优化内存使用的思路,比如在某些情况下,使用 Arrow IPC 文件可能比压缩的 Parquet 文件更合适。正如我之前所说的,务必根据手头的问题选择合适的格式,考虑所有因素:硬件、网络、数据以及整个流程。

下一章的标题是 跨越语言障碍:Arrow C 数据 API。到目前为止,我们已经覆盖了很多基础内容:创建数组和表格、读写各种格式的文件等。但记住,我提到过 Arrow 是一系列用于处理数据的库。接下来,我们将深入探索,研究 Arrow 库中一些更具实验性的 API,例如计算和数据集 API。本章的代码较少,但下一章不会如此。我希望你继续关注!

我们下章见!