使用 Apache Arrow 进行内存分析——使用关键的 Arrow 规范

418 阅读58分钟

分析和计算工具只有在拥有数据的情况下才有用。数据可以存在于本地或远程的多种位置和格式中。Arrow 库提供了一系列功能,用于从不同的格式和位置读取数据并进行交互。在你已经对 Arrow 及其数组操作有了扎实理解的基础上,本章将教你如何将数据导入到 Arrow 格式中,并在不同进程之间进行通信。

在本章中,我们将涵盖以下主题:

  • 从多种格式(包括 CSV、Apache Parquet 以及 pandas 和 Polars 的 DataFrame)导入数据
  • Arrow 与 pandas 数据和 Polars 数据的交互
  • 利用共享内存实现近乎零成本的数据共享

技术要求

在本章中,我将提供各种代码示例,使用 Python、C++ 和 Golang 的 Arrow 库以及公开的纽约出租车行程持续时间数据集(NYC Taxi Trip Duration Dataset)。

要运行本章中的实践示例,你需要以下环境:

  • 一台连接互联网的计算机
  • 你偏好的 IDE(VS Code、Sublime、Emacs、Vim 等)
  • 安装了 pyarrowpandaspolars 模块的 Python 3.8+
  • Go 1.21+
  • 带有 Arrow 库的 C++ 编译器(支持编译 C++17 或更新版本)
  • 本书 GitHub 仓库中的 sample_data 文件夹:github.com/PacktPublis…
  • 安装 git-lfs 以下载示例数据;本章中提供了安装它的说明

让我们开始吧!

随时随地玩转数据!

现代数据科学、机器学习(ML)以及其他数据操作技术通常需要从多个位置合并数据以执行任务。通常,这些数据并不在本地,而是存储在某种云存储中。大多数 Arrow 库的实现原生支持本地文件系统访问、Amazon Web Services 的简单存储服务(AWS S3)、Microsoft Azure 文件系统、Google Cloud Storage(GCS)和 Hadoop 分布式文件系统(HDFS)。除了原生支持的系统之外,文件系统接口通常在语言特定的情况下实现或使用,以便轻松添加对其他文件系统的支持。

一旦你能够访问存放文件的平台(无论是本地、云端或其他位置),你需要确保数据是 Arrow 库支持导入的格式。请查看你喜欢的编程语言的 Arrow 库文档,以了解支持的数据格式。Arrow 库提供的抽象使你能够轻松创建一个操作数据的流程,无论数据的位置或格式如何都能正常工作,然后将其写出到你希望的不同格式中。下图仅显示了大多数 Arrow 库支持的一些数据格式,但请记住——一个格式未列出并不一定意味着它不被支持:

image.png

在上图中,虚线轮廓显示了数据文件在处理之前可能存在于某个位置(例如 S3 或 HDFS 集群中);而处理的结果可以写入完全不同的存储位置和格式中。

为了提供优化和一致的库函数使用体验,以及简化实现,许多 Arrow 库定义了特定的文件系统使用接口。这些接口的具体实现方式会因语言而异,但它们的共同目标是抽象文件系统与导入数据文件交互的细节。在我们直接处理数据文件之前,需要引入几个重要的 Arrow 概念。我们已经在第 1 章《Apache Arrow 入门》中讨论了 Arrow 数组和记录批次,现在我们来介绍分块数组和表(chunked arrays 和 tables)。

使用 Arrow 表

简要回顾一下,记录批次是由相同长度的 Arrow 数组和描述列的名称、类型和元数据的模式组成的集合。通常,在读取和操作数据时,我们会以块的形式获取这些数据,然后希望将其组合起来,当作一个大表格进行处理,如下图所示:

image.png

一种方法是简单地分配足够的空间来存放整个表格,然后将每个记录批次的列复制到已分配的空间中。这样,我们就可以将最终的表格作为一个完整的记录批次保存在内存中。然而,这种方法有两个很大的问题,使得它无法扩展:

  1. 为每列分配一个全新的大块内存并复制所有数据的成本可能非常高。
  2. 如果我们再获取一个新的记录批次怎么办?每次获取更多数据时,我们都必须重新执行这个过程以适应现在变得更大的表格。

这时,分块数组(chunked arrays)的概念就派上了用场,如下图所示:

image.png

分块数组(Chunked Array)只是对一组具有相同数据类型的 Arrow 数组的一个轻量包装。这样,我们可以逐步高效地构建一个数组,甚至是一个完整的表格,而不必不断分配越来越大的内存块并复制数据。同样,Arrow 表保存了一个或多个分块数组和一个模式,类似于记录批次保存常规 Arrow 数组和模式的方式。表允许我们在概念上将所有数据视为一个连续的表,而无需频繁重新分配和复制数据。

这种方法的另一个好处是,不需要强制对齐列的块大小。例如,Field 0 可以分成 100 个元素的块,而 Field 1 可以分成 250 个元素的块,以此类推;核心点在于,这种方式使得列(进而是表格)可以以最少的数据复制次数构建,同时保持易于操作。

当然,这种方式也有一些权衡:由于数组不再是完全连续的缓冲区,我们会失去一些内存局部性。为了在处理数据时尽可能获得局部性带来的好处,块的大小应该尽可能大,这就需要在分配和复制的开销与处理非连续数据的开销之间进行平衡。幸运的是,大多数这种复杂性已经由 Arrow 库在输入/输出(I/O)接口中处理好了,因此在读取数据时可以进行高效的处理。然而,理解这些概念对于最大化数据集和操作的性能至关重要。

讲解完这些概念后,现在我们开始读取和写入一些文件吧!

为了简洁起见,这一部分我们将重点关注 Python 和 C++。但不用担心!Golang 在后续的示例中还会出现。在接下来的部分中,我们将看看如何利用可用的文件系统接口,从不同的支持文件格式中导入数据。首先是 Python!

使用 PyArrow 访问数据文件

Python Arrow 库定义了一个基础类接口,并针对不同的文件位置提供了一些具体实现,以便访问文件,如下图所示:

image.png

抽象接口 FileSystem 提供了用于 I/O 流和目录操作的实用工具。通过抽象底层文件系统交互的实现,提供了一个简化底层数据存储的单一接口。无论底层操作系统如何,路径始终用正斜杠(/)分隔,不包含特殊路径组件(如 ...),仅暴露文件的基本元数据,如大小和最后修改时间。构建 FileSystem 对象时,你可以显式地构建所需类型,或者从 URI 推断,例如:

  • 本地文件系统LocalFileSystem 构造函数接受一个可选参数 use_mmap。它默认是 False,如果设置为 True,则打开文件时会使用内存映射。有关内存映射文件的详细说明将在第 3 章“格式和内存处理”的“学习内存制图”部分中介绍。我们来构建这个对象:
>>> from pyarrow import fs
>>> local = fs.LocalFileSystem()  # 创建本地文件系统实例
>>> f, p = fs.FileSystem.from_uri('file:///home/mtopol/')
>>> f
<pyarrow._fs.LocalFileSystem object at 0x0000021FAF8F6570>
>>> p
'/home/mtopol/'

标准的 Windows 路径(如 C:\Users\mtopol...)无法使用,因为它们包含冒号。相反,你可以使用带正斜杠的 URI 指定路径:file:///c/Users/mtopol/...

  • AWS S3S3FileSystem 构造函数允许你通过多个参数指定凭证和其他连接属性(如区域或端点重写)。另外,构造函数还会检查标准的 S3 凭证配置(例如,环境变量 AWS_ACCESS_KEY_ID~/.aws/config 文件,或对于 Amazon EC2 上的节点,检查 EC2 实例的元数据服务)。如果你想以匿名方式连接(无需凭证),还可以使用一个名为 anonymous 的布尔参数来表示:
>>> from pyarrow import fs
>>> s3 = fs.S3FileSystem(region='us-east-1')  # 显式创建
>>> s3, path = fs.FileSystem.from_uri('s3://my-bucket/')  # 从 URI 推断
>>> s3
<pyarrow._s3fs.S3FileSystem object at 0x0000021FAF7F99F0>
>>> path
'my-bucket'
  • GCSGcsFileSystem 构造函数允许你指定以匿名方式连接(无需凭证),或直接提供访问令牌。如果两者都未指定,则将根据 Google 文档中描述的流程(google.aip.dev/auth/4110)来检测和选择凭证。如果你不在 Google Cloud Platform(GCP)上运行,这通常需要设置环境变量 $GOOGLE_APPLICATION_CREDENTIALS 指向包含所需凭证的 JSON 文件:
>>> from datetime import timedelta
>>> from pyarrow import fs
>>> gcs = fs.GcsFileSystem(anonymous=True)  # 显式构建
>>> gcs, p = fs.GcsFileSystem.from_uri('gs://anonymous@gcp-public-data-landsat')  # 从 URI 构建
>>> gcs
<pyarrow._gcsfs.GcsFileSystem object at 0x7f338f6f1b30>
>>> uri = "gcp-public-data-landsat/LC08/01/001/003/"
>>> file_list = gcs.get_file_info(fs.FileSelector(uri, recursive=True))
>>> len(file_list)
814
  • HDFS:使用 HDFS 稍微有点复杂,因为它需要在你的路径上包含 Java Native Interface (JNI) 库,以便它们可以加载。JNI 是一个框架,允许运行在 Java 虚拟机(JVM)中的 Java 代码调用并被本地应用程序和库调用。这里我不会详细介绍 HDFS 的安装(你可以在 Hadoop 文档中找到相关内容:hadoop.apache.org/docs/curren…),但要与 PyArrow 一起使用 HDFS 的关键是确保 libjvm.solibhdfs.so$LD_LIBRARY_PATH 中(在 Windows 上为 $PATH,在新版 macOS 上为 $DYLD_FALLBACK_LIBRARY_PATH),以便它们能够在运行时加载。如果这些库在运行时可访问,你可以使用 PyArrow 与 HDFS 集群通信:
>>> from pyarrow import fs
>>> hdfs = fs.HadoopFileSystem(host='namenode', port=8020)
>>> hdfs, path = fs.FileSystem.from_uri('hdfs://namenode:8020/tmp')
>>> hdfs
<pyarrow._hdfs.HadoopFileSystem object at 0x7f7a70960bf0>
>>> path
'/tmp'

PyArrow 库将在构建时尝试连接到 HDFS 的名称节点,如果连接不成功则会失败。Hadoop 库的运行时查找依赖于几个不同的环境变量。如果库不在 $LD_LIBRARY_PATH 环境变量中,你可以使用以下环境变量配置查找路径。

  • 如果你有完整的 Hadoop 安装,应该定义 HADOOP_HOME,它通常包含 lib/native/libhdfs.soJAVA_HOME 环境变量应定义为指向 Java SDK 的安装路径。如果 libhdfs.so 安装在 $HADOOP_HOME/lib/native 之外的某处,你可以使用 ARROW_LIBHDFS_DIR 环境变量指定其位置。

PyArrow 中的许多与 I/O 相关的函数允许调用方指定一个 URI 以推断文件系统,或者提供一个明确的参数来指定要使用的 FileSystem 实例。一旦初始化了所需的文件系统实例,就可以使用接口进行许多标准文件系统操作,无论底层实现如何。以下是一些抽象函数的子集,帮助你快速入门:

  • create_dir:创建目录或子目录
  • delete_dir:递归删除目录及其内容
  • delete_dir_contents:递归删除目录内容
  • copy_filedelete_file:按路径复制或删除指定文件
  • open_input_file:打开文件进行随机访问读取
  • open_input_stream:打开文件,仅用于顺序读取
  • open_append_stream:打开输出流以追加数据
  • open_output_stream:打开输出流,用于顺序写入

除了操作文件和目录,你还可以使用抽象层检查和列出文件和目录的内容。打开文件或流会生成一个所谓的“文件对象”,该对象可以与任何支持此类对象的函数一起使用,无论底层存储或位置如何。

到这里为止,你应该已经掌握了如何使用 Python Arrow 库打开并引用数据文件。接下来,我们可以开始研究本地实现的不同数据格式,以及如何将这些数据格式处理成 Arrow 数组和表。

使用 PyArrow 处理 CSV 文件

用于数据的最常见文件格式之一是分隔文本文件,例如逗号分隔值 (CSV) 文件。除了逗号外,它们还常被用作制表符或竖线分隔文件。由于 CSV 文件的原始文本没有明确的类型定义,Arrow 库会尝试推断类型,并提供了多种选项来在读取或写入时解析和转换数据为 Arrow 数据。关于类型推断的更多信息,可以参考 Arrow 文档:arrow.apache.org/docs/python…

默认的 CSV 文件读取选项在推断数据类型方面通常效果很好,因此读取简单的文件非常容易。我们可以通过使用 train.csv 示例数据文件来查看这一点,它是常用的纽约出租车行程持续时间数据集的一个子集,这个数据集因常用于机器学习训练数据集而得名:

>>> import pyarrow as pa
>>> import pyarrow.csv
>>> table = pa.csv.read_csv('sample_data/train.csv')
>>> table.schema
vendor_id: string
pickup_at: timestamp[ns]
dropoff_at: timestamp[ns]
passenger_count: int64
trip_distance: double
pickup_longitude: double
pickup_latitude: double
rate_code_id: int64
store_and_fwd_flag: string
dropoff_longitude: double
dropoff_latitude: double
payment_type: string
fare_amount: double
extra: double
mta_tax: double
tip_amount: double
tolls_amount: double
total_amount: double

首先需要注意的是,我们直接将字符串传递给 read_csv 函数。与许多 Python Arrow 文件读取函数一样,它可以接受以下参数之一:

  • 字符串:文件名或文件路径。如果有必要,将推断文件系统。
  • 文件对象:这可以是内置 open 函数返回的对象,或者实现了 FileSystem 接口并返回可读对象的函数,例如 open_input_stream

文件的第一行包含列标题,这些标题会自动用作生成的 Arrow 表的列名称。之后我们可以看到,库从值中识别出时间戳列,并确定精度为秒,而不是毫秒或纳秒。最后,可以看到数值列和字符串列,这些决定了相应列是双精度数而非整数列。

约定和假设

CSV 文件的第一行包含列标题,不幸的是,这只是我们遵循的一个惯例,而非必须要求。偏离这种惯例的情况非常常见,这使得处理 CSV 文件变得困难。但不必担心,之前提到的文档中也包含了如何处理不符合该惯例的文件的相关信息!

读取 CSV 文件会返回一个 pyarrow.Table 类型的对象,其中包含一组 pyarrow.lib.ChunkedArray 对象。这符合前面提到的关于表和分块数组的模式。在读取文件时,可以通过一次读取若干行并构建分块列来进行并行化,而无需复制数据。下图显示了并行化文件读取的过程:

image.png

在这里,我们可以看到线程正在并行读取文件。每个线程将一组行读入数组块中。一旦创建了 Arrow 数组,这些数组块会被零拷贝地添加到表的列中。我们可以在 Python 解释器中查看完成的表的列,以便看到这个过程的效果:

>>> table.column(0).num_chunks
192
>>> table.column(0).chunks
[<pyarrow.lib.StringArray object at 0x000001B2C5EB9FA8>[  "VTS",  "VTS",  "VTS",  "VTS",  "VTS",  "VTS",  "VTS",  ...  "CMT",  "VTS",  "VTS",  "VTS",  "VTS",  "VTS",  "CMT",  "VTS",  "VTS",  "VTS"], <pyarrow.lib.StringArray object at 0x000001B2C5EBB048>
[……

除了输入流或文件名之外,CSV 读取函数还有三种类型的选项可以传递进去——ReadOptionsParseOptionsConvertOptions。每个选项集允许控制读取文件和创建 Arrow 表对象的不同方面,具体如下:

  • ReadOptions:此选项允许你配置是否使用线程进行读取、每次读取的块大小,以及如何为表生成列名(可以从文件中读取、直接提供或自动生成)。(需要注意的是,如果文件的扩展名是已知的压缩格式,数据将在读取时自动解压缩,而不需要你自己解压缩!):
table = pa.csv.read_csv('file.csv.gz',
    read_options=pa.csv.ReadOptions(
        encoding='utf8', # 文件的编码类型
        column_names=['col1', 'col2', 'col3'],
        block_size=4096,  # 每次处理的字节数
)
  • ParseOptions:此选项控制用于分隔列的分隔符、引号和转义字符,以及如何处理换行符:
table = pa.csv.read_csv(input_file, parse_options=pa.csv.ParseOptions(
    delimiter='|',  # 适用于竖线分隔文件
    escape_char='\',  # 允许反斜杠转义值
)
  • ConvertOptions:此选项提供了将数据转换为 Arrow 数组数据的各种选项,包括指定哪些字符串应被视为 null 值、哪些字符串应被视为布尔值中的 truefalse,以及其他各种将字符串解析为 Arrow 数据类型的选项:
table = pa.csv.read_csv('tips.csv',
    convert_options=pa.csv.ConvertOptions(
       column_types={
           'total_bill': pa.decimal128(precision=10, scale=2),
           'tip': pa.decimal128(precision=10, scale=2),
       },
       # 仅从文件中读取这些列,并按此顺序读取
       # 忽略其他列
       include_columns=['tip', 'total_bill', 'timestamp'],
))

除了读取 CSV 文件的所有功能外,还有一个 write_csv 函数可以从记录批或表中写入 CSV 文件。与 read 函数类似,它接受文件名或路径,或一个可以写入的文件对象作为参数。写入时,有几个可用选项可以操纵写入的内容:

  • 是否包含列名称的初始标题行
  • 写出行时使用的批处理大小
  • 选择并指定文件写出的分隔符和引用样式

下面是一个简单的函数示例,它可以从 CSV 文件中读取某些列的子集,并将该子集写入一个新文件:

def create_subset_csv(input, output, column_names):
    table = pa.csv.read_csv(input,
        convert_options=pa.csv.ConvertOptions(
            include_columns=column_names))
    pa.csv.write_csv(table, output,
                     write_options=pa.csv.WriteOptions(
                         include_header=True))

在某些情况下,你可能希望在生成或检索数据时,将数据增量地写入 CSV 文件。当你这样做时,如果可以避免,你不希望将整个表一次性加载到内存中。此时,你可以使用 pyarrow.csv.CSVWriter 来增量地写入数据:

schema = pa.schema([("col", pa.int64())])
with pa.csv.CSVWriter("output.csv", schema=schema) as writer:
    for chunk in range(10):
        datachunk = range(chunk*10, (chunk+1)*10)
        table = pa.Table.from_arrays([pa.array(datachunk)], schema=schema)
        writer.write(table)

接下来我们要介绍的另一种非常常见的数据格式是 JSON 数据。

使用 PyArrow 处理 JSON 文件

JSON 数据文件的预期格式是每行包含一个 JSON 对象,表示一行数据的换行分隔文件。读取 JSON 文件的过程与读取 CSV 文件几乎相同!以下是一个示例 JSON 数据文件:

{"a": 1, "b": 2.0, "c": 1}
{"a": 3, "b": 3.0, "c": 2}
{"a": 5, "b": 4.0, "c": 3}
{"a": 7, "b": 5.0, "c": 4}

将该文件读取到表中非常简单:

>>> import pyarrow as pa
>>> import pyarrow.json
>>> table = pa.json.read_json(filename)
>>> table.to_pydict()
{'a': [1, 3, 5, 7], 'b': [2.0, 3.0, 4.0, 5.0], 'c': [1, 2, 3, 4]}

与读取 CSV 文件类似,读取 JSON 文件也有 ReadOptionsParseOptions 选项,可让你配置创建 Arrow 数据的行为。你可以通过指定明确的架构和定义如何处理意外字段来进行操作。当前,还没有相应的 write_json 函数。

由于 JSON 有一些类型规则,如果未提供架构,它在推断数据类型时所需的工作量会少一些。推断类型的规则如下:

  • JSON 中的 null 值会转换为 null 数据类型,但在该列中遇到非 null 值时会回退到其他类型。
  • 布尔值会被转换为 Arrow 的布尔数据类型。
  • 数字数据会被推断为 64 位整数,如果找到非整数值,则回退为 64 位浮点数。
  • 格式为 "YYYY-MM-DD" 和/或 "YYYY-MM-DD hh:mm:ss" 的字符串将被解释为时间戳数据类型,单位为秒。如果发现任何错误,将回退为常规字符串类型。
  • 数组会被推断为列表类型,递归推断列表中元素的类型。
  • 嵌套对象会被推断为结构数据类型,其字段会根据这些规则递归推断。

更多信息可以参考 PyArrow 的 JSON API 文档:arrow.apache.org/docs/python…

使用 PyArrow 处理 ORC 文件

与 JSON 和 CSV 格式不同,Apache 优化的列式存储格式(ORC)不太流行,除非你已经在使用 Hadoop 数据生态系统。Apache ORC 是一种列式格式,最初由 Hortonworks 开发,用于以压缩格式存储数据,以便通过 Apache Hive 和其他 Hadoop 工具进行处理。它以列为单位存储数据,配合文件索引,并将文件分割成条带(stripes),以促进谓词下推和优化读取。

由于 ORC 文件格式经常用于数据存储和查询,PyArrow 库提供了一个接口,可以直接将 ORC 文件读取为 Arrow 表,该接口名为 pyarrow.orc.ORCFile。与 Arrow 类似,ORC 文件具有模式(schema),并且列具有明确的类型,因此它们可以轻松转换。因为不再像从 JSON 或 CSV 文件中推断数据类型那样存在不确定性。我们可以调整之前的示例,让它读取 ORC 文件而不是 CSV 或 JSON 文件:

要将 ORC 文件读取为一个或多个 Arrow 表,首先创建一个 ORCFile 实例:

>>> import pyarrow as pa
>>> import pyarrow.orc
>>> of = pa.orc.ORCFile('train.orc')
>>> of.nrows
1458644
>>> of.schema
vendor_id: string
pickup_at: timestamp[ns]
dropoff_at: timestamp[ns]
passenger_count: int64
trip_distance: double
pickup_longitude: double
pickup_latitude: double
rate_code_id: int64
store_and_fwd_flag: string
dropoff_longitude: double
dropoff_latitude: double
payment_type: string
fare_amount: double
extra: double
mta_tax: double
tip_amount: double
tolls_amount: double
total_amount: double

与读取 CSV 和 JSON 文件类似,创建 ORCFile 实例的参数可以是文件路径或文件对象。

有了这个对象,你现在可以将整个文件或其中的一个条带(stripe)读取为 Arrow 表,甚至可以只读取所需的部分列:

>>> tbl = of.read(columns=['vendor_id', 'passenger_count', 'rate_code_id'])
# 留空或使用 None 来读取所有列
>>> tbl
pyarrow.Table
vendor_id: string
passenger_count: int64
rate_code_id: int64
---
id: [["VTS","VTS","VTS","VTS","VTS",]]
passenger_count: [
    [1,1,2,1,1,1,1,1,1,2,...,5,5,1,1,1,3,5,4,1,1]]
trip_duration: [[1,1,1,1,1,,1,1,1,1,1,1,1]]

除了读取 ORC 文件外,我们还可以使用 pyarrow.orc.write_table 写入 ORC 文件。此方法的参数是要写入的表和文件的目标位置。

使用 PyArrow 处理 Apache Parquet 文件

如果你不熟悉 Parquet,它与 ORC 类似,都是列式的磁盘存储格式,并带有压缩功能。两者都包含各种元数据,以便直接从文件中进行高效查询。可以将它们看作是两种不同风格的列式存储格式,它们在设计中做出了不同的权衡。

注意
对于所有这些列式存储格式,你可能会想知道为什么还需要 Arrow,什么时候使用哪种格式,出于什么原因?别担心!我们将在第 3 章 "格式和内存处理" 中详细探讨这些问题,以及其他格式比较的问题。

到了这一步,你可能已经发现了这个库设计中的模式,并且可以猜到将 Parquet 文件读取为 Arrow 表的代码可能是什么样子。试试看,自己写一个代码草稿;我等你。

想好了吗?好吧,让我们来看看:

>>> import pyarrow.parquet as pq
>>> table = pq.read_table('train.parquet')

没错,就这么简单。当然,有很多选项可以自定义和优化 Parquet 文件的读取。可用的一些选项如下:

  • 指定要读取的列列表,以便只读取包含这些列数据的文件部分。
  • 控制缓冲区大小,以及是否预先将文件中的数据缓冲到内存中以优化 I/O 操作。
  • 可以传入一个文件系统选项,使得文件路径使用提供的 FileSystem 对象进行查找和打开,而不是本地磁盘系统。
  • 通过过滤选项推送谓词以过滤掉行数据。

Python 部分介绍完毕,接下来我们看看 C++ 库如何处理相同的功能和连接操作。

使用 C++ 访问 Arrow 数据文件

由于 Python 接口是构建在 C++ 库之上的,因此接口非常类似于 PyArrow 库的 fs 模块。你可以通过在代码中包含 arrow/filesystem/api.h 来使用 C++ 库的文件系统模块,它将加载 arrow::fs 命名空间中的主要文件系统处理程序 —— LocalFileSystem、S3FileSystem、GcsFileSystem、AzureFileSystem 和 HadoopFileSystem,这些都是与 Python 库中相同的具体实现。它们都提供了在其各自物理位置上创建、复制、移动和读取文件的基本功能,巧妙地抽象出复杂性以简化使用。当然,正如 Python 模块中的 from_uri 函数一样,C++ 中也有 arrow::fs::FileSystemFromUriarrow::fs::FileSystemFromUriOrPath,它们可以从提供的 URI 或本地文件路径构建文件系统实例。

接下来,让我们通过一些示例来展示如何使用这些功能与各种数据格式进行交互。我们将从 CSV 文件开始。

使用 Arrow 处理 C++ 中的 CSV 数据文件

默认情况下,Arrow 库会读取 CSV 文件中的所有列。不过,提供了多种选项来控制如何处理文件。以下是一些选项的示例,建议你查阅文档以获取完整选项列表:

  • 默认情况下,列名将从 CSV 文件的第一行读取;否则,可以使用 arrow::csv::ReadOptions::column_names 设置列名。如果设置了该选项,则文件的第一行将作为数据读取。
  • arrow::csv::ConvertOptions::include_columns 可用于指定要读取的列,并忽略其他列。除非将 ConvertOptions::include_missing_columns 设置为 false,否则如果缺少所需的列,将返回错误;否则,缺少的列将以全空值的形式返回。
  • CSV 读取器默认会推断列的数据类型,但可以通过可选的 arrow::csv::ConvertOptions::column_types 映射来指定数据类型。
  • arrow::csv::ParseOptions 包含各种字段,用于自定义如何将文本解析为值,例如指示 truefalse 的值、列分隔符等。

现在,让我们看一个代码示例:

首先,我们需要包含以下头文件:

#include <arrow/io/api.h>  // 用于打开文件
#include <arrow/csv/api.h> // CSV 功能和对象
#include <arrow/table.h>   // 因为我们将数据读取为表格
#include <iostream>        // 用于在终端输出

接下来,我们将打开文件。为了简化起见,我们暂时使用本地文件:

auto maybe_input = arrow::io::ReadableFile::Open("train.csv");
if (!maybe_input.ok()) {
    // 处理文件打开错误
    maybe_input.status()
}
std::shared_ptr<arrow::io::InputStream> input = *maybe_input;

然后,我们需要设置选项对象。我这里将使用默认选项,但你可以尝试不同选项组合:

auto io_context = arrow::io::default_io_context();
auto read_options = arrow::csv::ReadOptions::Defaults();
auto parse_options = arrow::csv::ParseOptions::Defaults();
auto convert_options = arrow::csv::ConvertOptions::Defaults();

现在,一切准备就绪,我们只需创建表格读取器并获取数据:

auto maybe_reader = arrow::csv::TableReader::Make(io_context, input, read_options, parse_options, convert_options);
if (!maybe_reader.ok()) {
    // 处理 TableReader 实例化错误
}
std::shared_ptr<arrow::csv::TableReader> reader = *maybe_reader;

// 从文件中读取数据表
auto maybe_table = reader->Read();
if (!maybe_table.ok()) {
    // 处理错误,例如 CSV 语法错误或类型转换失败等
}
std::shared_ptr<arrow::Table> table = *maybe_table;

:你可能注意到函数返回 arrow::Result 对象,这些对象是模板化的值。这样可以方便检查处理过程中是否有错误,并优雅地处理,而不是在运行时崩溃。我们使用 ok 方法检查是否成功,并通过 status 方法获取错误代码和消息。

现在我们已经将整个文件读入内存(因为这是默认选项的行为),我们可以将其打印出来:

std::cout << table->ToString() << std::endl;

完整的代码示例可以在本书的 GitHub 仓库的 chapter2 目录下找到,文件名为 csv_reader.cc。编译时,请确保正确链接库。执行后,输出结果将与我们在 Python 中使用 PyArrow 时几乎相同。与之前一样,数据被读入到表的分块数组中,以便在读取操作期间进行并行处理。尝试使用不同的选项来读取不同的表结果,并控制读取到内存中的内容。

写入 CSV 文件同样非常简单。与读取操作一样,你可以一次性写入整个表,或进行增量写入。完整代码示例可以在本书 GitHub 仓库的 chapter2 目录下找到,文件名为 csv_writer.cc

arrow::Table table = …;
// 一次性写入表格
bool append = false; // 设置为 true 以附加到现有文件
auto maybe_output = arrow::io::FileOutputStream::Open("train.csv", append);
if (!maybe_output.ok()) {
    // 处理错误
}
auto output = *maybe_output;
auto write_options = arrow::csv::WriteOptions::Defaults();
auto status = arrow::csv::WriteCSV(*table, write_options, output.get());
if (!status.ok()) {
    // 处理错误,打印 status.message()
}

要进行增量写入,你需要创建一个 CSVWriter 对象并逐步将记录批次写入文件。

使用 Arrow 在 C++ 中处理 JSON 数据文件

JSON 文件的预期格式与 Python 库相同:以换行符分隔的 JSON 对象,每个输入文件中的对象对应于生成的 Arrow 表中的一行。语义上,读取 JSON 数据文件的方式与 Python 类似,提供选项来控制数据的转换,或者让库自动推断数据类型。在 C++ 中处理 JSON 文件与我们之前处理 CSV 文件非常相似。

读取 JSON 文件与 CSV 文件的主要区别如下:

  1. 包含 <arrow/json/api.h> 而不是 <arrow/csv/api.h>
  2. 调用 arrow::json::TableReader::Make 而不是 arrow::csv::TableReader::Make。JSON 读取器接受 MemoryPool* 对象,而不是 IOContext 对象,其他工作方式类似。

现在我们已经在 C++ 中处理了基于文本的、可人类读取的文件格式,接下来让我们转向二进制格式。按照与 Python 相同的顺序,接下来让我们尝试使用 C++ 库读取 ORC 文件。

使用 Arrow 在 C++ 中处理 ORC 数据文件

对 ORC 的支持由官方的 Apache ORC 库(即 liborc)提供。如果 Arrow 库是从源代码编译的,则可以选择是否构建 ORC 支持。然而,官方发布的 Arrow 包应该都已经内置了 ORC 适配器,前提是有官方的 ORC 库。

与 CSV 或 JSON 读取器不同的是,写作时的 ORC 读取器不支持流,它只能使用 arrow::io::RandomAccessFile 的实例。幸运的是,打开文件系统中的文件会生成这样的类型,因此它不会改变我们使用的基本模式。虽然这里的示例使用本地文件系统,但你始终可以通过使用各自的文件系统抽象实例化与 S3 或 HDFS 集群的连接,并以相同的方式从它们打开文件。

让我们开始吧:

在打开文件后,我们需要创建 ORCFileReader 读取器:

#include <arrow/adapters/orc/adapter.h>

// 不显式处理错误,而是使用 ValueOrDie 抛出异常
std::shared_ptr<arrow::io::RandomAccessFile> file =
    arrow::io::ReadableFile::Open("train.orc").ValueOrDie();
arrow::MemoryPool* pool = arrow::default_memory_pool();
auto reader = arrow::adapters::orc::ORCFileReader::Open(file, pool).ValueOrDie();

ORC 读取器提供了许多不同的函数,例如读取特定的条带(stripes),查找特定的行号,检索条带或行的数量等。不过现在我们只是将整个文件读入一个表:

std::shared_ptr<arrow::Table> data = reader->Read().ValueOrDie();

最后,我们可以以相同的方式写入 ORC 文件:

std::shared_ptr<arrow::io::OutputStream> output =
    arrow::io::FileOutputStream::Open("train.orc").ValueOrDie();
auto writer = arrow::adapters::orc::ORCFileWriter::Open(output.get()).ValueOrDie();
status = writer->Write(*data);
if (!status.ok()) {
    // 处理写入错误
}
status = writer->Close();
if (!status.ok()) {
    // 处理关闭错误
}

由于 ORC 具有定义好的架构和类型支持,因此 Arrow 与 ORC 之间的转换非常清晰。可以直接从 ORC 文件中读取架构并转换为 Arrow 架构。为了优化读取模式,还可以选择只读取特定的条带。随着社区的发展,更多关于使用 ORC 文件与 Arrow 结合的功能将会被开发出来。

使用 Arrow 在 C++ 中处理 Parquet 数据文件

Parquet C++ 项目在不久前被并入 Apache Arrow 项目,因此它包含了许多功能,并且与 Arrow C++ 实用工具和类集成得非常好。我们不会在这里详细介绍 Parquet 的所有功能,不过它们非常值得一看,许多功能也会在后续章节中涉及或提到。

我们接下来要做的,是将 Parquet 文件读入内存中的 Arrow 表。我们将按照与 ORC 文件读取器相同的模式进行——只不过使用 Parquet Arrow 读取器。与 ORC 文件一样,我们需要一个 arrow::io::RandomAccessFile 实例作为输入,因为 Parquet 文件的元数据位于文件末尾的页脚(footer),描述了文件中每列数据的读取位置。

让我们开始吧:

我们要使用的 include 指令也会改变,创建输入文件实例后,我们可以创建一个 Parquet 读取器:

#include <parquet/arrow/reader.h>

std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
arrow::Status st = parquet::arrow::OpenFile(input, pool, &arrow_reader);
if (!st.ok()) {
    // 处理错误
}

你可以使用多种函数来控制只读取特定的行组、获取记录批次读取器、只读取特定列,等等。你甚至可以访问底层的 ParquetFileReader 对象(可在文档中找到大多数功能)。默认情况下,读取时只会使用一个线程,但你可以启用多线程来读取多个列。让我们来处理最简单的情况:

std::shared_ptr<arrow::Table> table;
st = arrow_reader->ReadTable(&table);
if (!st.ok()) {
    // 处理读取错误
}

从 Arrow 表写入 Parquet 文件的数据操作正如你现在所预期的那样进行:

#include <parquet/arrow/writer.h>

PARQUET_ASSIGN_OR_THROW(auto outfile, arrow::io::FileOutputStream::Open("train.parquet"));
int64_t chunk_size = parquet::DEFAULT_MAX_ROW_GROUP_LENGTH;
PARQUET_THROW_NOT_OK(parquet::arrow::WriteTable(table, arrow::default_memory_pool(), outfile, chunk_size));

快速提示:你可能会注意到 PARQUET_ASSIGN_OR_THROWPARQUET_THROW_NOT_OK 宏。它们相当于类似名称的 Arrow 宏,在 Arrow 和 Parquet 代码库中广泛使用。这些 *_ASSIGN_OR_* 宏会从 arrow::Result 对象中赋值,如果有效则赋值,否则抛出异常或返回错误状态。

现在你已经了解了如何在 C++ 中使用 Arrow 处理 JSON、ORC 和 Parquet 文件的数据。

熊熊们也会发射箭矢

如果你曾在 Python 中进行过数据分析,你很可能至少听说过 pandas 库。pandas 是一个开源的、BSD 许可的 Python 数据分析库,是数据科学家和工程师使用最广泛的工具之一。考虑到它的广泛使用,Arrow 的 Python 库为高效、快速地在 pandas DataFrame 和 Arrow 之间进行转换提供了集成功能。最近,Polars 库的开发也在 Rust 语言中使用 Arrow 作为其内部内存模型。

本节首先将深入探讨如何结合使用 Arrow 和 pandas 来提升工作流程的具体事项及注意事项。随后,我们还将介绍如何使用 Polars 并与 Arrow 数据共享。

在开始之前,确保你已经在本地安装了 pandas 和 Polars 以便跟随学习。当然,你还需要安装 PyArrow,这已经在上一章中提到过了。让我们来看一下安装步骤:

如果你使用的是 conda,pandas 和 Polars 都包含在 Anaconda(docs.continuum.io/anaconda/)发…

conda install pandas polars

如果你更喜欢使用 pip,可以通过 PyPI 正常安装:

pip3 install pandas polars

然而,情况并非总是如此顺利。目前还不可能将每个列类型都未修改地转换为 pandas。我们首先需要了解 Arrow 和 pandas 之间的数据类型是如何比较和对应的。

将 pandas 融入你的箭袋

pandas 的标准构建块是 DataFrame,它大致相当于 Arrow 表。两者都描述了具有相同长度的命名列。在最简单的情况下,Arrow 中有方便的 to_pandasfrom_pandas 函数用于转换。例如,表与 DataFrame 之间的转换非常简单:

>>> import pyarrow as pa
>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3]})
>>> table = pa.Table.from_pandas(df) # 转换为 Arrow 表
>>> df_new = table.to_pandas() # 转换回 pandas

此外,还有许多选项可以控制转换过程,例如是否使用线程、内存管理以及数据类型管理。pandas 对象还可能有一个 index 成员变量,它可以包含数据的行标签,而不仅仅使用 0 为基础的行索引。在从 DataFrame 进行转换时,from_pandas 系列函数有一个名为 preserve_index 的选项,用于控制是否以及如何存储索引数据。通常,索引数据会作为结果 Arrow 表的模式元数据进行跟踪。选项如下:

  • None:这是 preserve_index 选项的默认值。RangeIndex 类型的数据仅作为元数据存储,而其他索引类型的数据会作为实际列存储在生成的表中。
  • False:完全不存储任何索引信息。
  • True:强制将所有索引数据序列化为表中的列,因为存储 RangeIndex 实例可能在某些场景(例如将多个 DataFrame 存储在单个 Parquet 文件中)中引发问题。

Arrow 表支持扁平和嵌套列,而 DataFrame 仅支持扁平列。这种数据结构差异以及数据类型处理方式的差异,意味着有时无法进行完整的转换。主要的转换难点之一在于 pandas 不支持任意类型的可空列,而 Arrow 支持。所有 Arrow 数组都可能包含空值,无论其类型如何。另外,pandas 的日期时间处理始终使用纳秒作为时间单位。下图展示了数据类型之间的映射关系:

image.png

在 pandas 中,只有某些数据类型支持处理缺失数据或空值。一个特定的例子是默认的整数类型不支持空值。如果数组中有空值,从 Arrow 转换时,这些整数类型将会被强制转换为浮点类型;如果没有空值,列则保留其原有的整数类型,如以下代码片段所示:

>>> arr = pa.array([1, 2, 3])
>>> arr
<pyarrow.lib.Int64Array object at 0x000002348DCD02E8>
[
  1,
  2,
  3
]
>>> arr.to_pandas()
0  1
1  2
2  3
dtype: int64
>>> arr = pa.array([1, 2, None])
>>> arr.to_pandas()
0  1.0
1  2.0
2  NaN
dtype: float64

当你处理日期和时间类型时,有一些需要注意的事项,尤其是需要了解你所期望的数据类型:

  • 通常情况下,pandas 中的日期是通过 numpy.datetime64[ns] 类型处理的,但有时日期会表示为 Python 内置的 datetime.date 对象数组。默认情况下,两者都会被转换为 Arrow 的 date32 类型。
  • 如果你想使用 date64,则必须明确指定:
>>> from datetime import date
>>> s = pd.Series([date(1987, 8, 4), None, date(2000, 1, 1)])
>>> arr = pa.array(s)
>>> arr.type
DataType(date32[day])
>>> arr = pa.array(s, type='date64')
>>> arr.type
DataType(date64[ms])

转换回 pandas 时,你将得到 datetime.date 对象。使用 date_as_object=False 参数可以获取 NumPy 的 datetime64 类型:

>>> arr.to_pandas()
0   1987-08-04
1   None
2   2000-01-01
dtype: object
>>> s2 = pd.Series(arr.to_pandas(date_as_object=False))
>>> s2.dtype
dtype('<M8[ns]')

Python 内置的 datetime.time 对象将会被转换为 Arrow 的 time64 类型,反之亦然。pandas 使用的时间戳类型是 NumPy 的 datetime64[ns],这会被转换为 Arrow 中使用纳秒为单位的时间戳类型。如果使用时区信息,它会在转换时通过元数据保留。此外,注意到基础值在 Arrow 中被转换为 UTC,并且时区信息会存储在数据类型的元数据中:

>>> df = pd.DataFrame({'datetime': pd.date_range('2020-01-01T00:00:00-04:00', freq='H', periods=3)})
>>> df
                   datetime
0 2020-01-01 00:00:00-04:00
1 2020-01-01 01:00:00-04:00
2 2020-01-01 02:00:00-04:00
>>> table = pa.Table.from_pandas(df)
>>> table
pyarrow.Table
datetime: timestamp[ns, tz=-4:00]
datetime: [[2020-01-01 04:00:00.000000000,2020-01-01 05:00:00.000000000,2020-01-01 06:00:00.000000000]]

由于 pandas 在 Python 生态系统中使用广泛,并且已经提供了读写 CSV、Parquet 等文件类型的工具,因此除了互操作性之外,Arrow 所提供的优势可能并不那么明显。但不可否认,为 pandas 提供与任何基于 Arrow 的工具的互操作性已经是一个巨大的好处!特别是当你考虑内存使用以及读写和传输数据的性能时,Arrow 的优势更为突出。

进展继续!

从 pandas 2.0 开始,你可以直接让 pandas 使用 PyArrow 将其底层数据存储在 Arrow 数组中,而不是之前使用的传统 NumPy 后端。这改善了 pandas 和 PyArrow 之间的转换(更好的零拷贝支持!),并且证明了我们所强调的直接使用 Arrow 的优势。

让 pandas 跑得更快

借助 Arrow 确保高效内存使用和数据传输的所有优化和底层调整,我们有必要比较使用 Arrow 库和 pandas 库读取数据文件的速度。利用 IPython 工具,可以轻松进行时间测试比较。我们将使用前面示例中用于读取数据文件的相同示例数据文件来进行测试:

In [1]: import pyarrow as pa
In [2]: import pyarrow.csv
In [3]: %timeit table = pa.csv.read_csv('train.csv')
177 ms ± 3.03 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

上面的输出显示了使用 timeit 工具读取 CSV 文件时的结果,它使用 PyArrow 读取了七次,并得出了每次运行的平均值和标准偏差。在我的笔记本电脑上,从 CSV 文件创建 Arrow 表平均只需要 177 毫秒,该文件大约有 192 MB 大小。为了保持公平的比较,我们还需要测试从 Arrow 表创建 pandas DataFrame 所需的时间,以便做出真正的对比:

In [4]: table = pa.csv.read_csv('train.csv')
In [5]: import pandas as pd
In [6]: %timeit df = table.to_pandas()
509 ms ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

可以看到,大约需要 509 毫秒将表转换为 DataFrame,而这远远超过了将文件读取到 Arrow 表中所需的时间。现在,让我们看看使用 pandas 的 read_csv 函数读取同一文件需要多长时间:

In [7]: %timeit df = pd.read_csv('train.csv')
3.49 s ± 193 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

哇!看看这个!在我的笔记本电脑上,直接使用 pandas 读取文件平均耗时 3.5 秒。即使结合将文件读入 Arrow 表和从 Arrow 表转换为 DataFrame 的时间,总共仍然比 pandas 快约 80%,即使这个文件在数据分析标准下算是相对较小的文件(包含 1,458,644 行和 11 列)。为了给 pandas 一个公平的机会,我们还可以尝试从压缩版的 CSV 文件中读取,增加了解压缩数据所需的处理步骤,然后再解析以创建最终的对象。以下图表显示了使用 timeit 工具的最终时间,不仅包括读取文件及其压缩形式的时间,还包括从数据写入 CSV 文件的时间:

image.png

你可能会好奇,除了 CSV 之外的其他文件格式,以及 pandas 和 PyArrow 之间的性能比较。如果你查看 pandas 处理 Parquet 和 ORC 文件格式的函数文档,你会发现它们实际上都调用了 PyArrow 库来读取数据。因此,在这两种情况下,pandas 只是将任务委托给了 PyArrow 来完成。而对于 JSON 文件,pandas 期望的数据结构和格式与 PyArrow 的不同,因此在这种情况下,它们并不具有等价性。你应当根据实际需求来选择具体的工具,这通常取决于你所使用的数据源格式。

有时,在 Arrow 数组或表与 pandas DataFrame 之间进行转换时,内存使用和性能问题可能会变得突出。因为这两个库对底层数据的内部表示不同,所以只有在有限的情况下,转换才能在不复制数据或不执行额外计算的前提下进行。在最糟糕的情况下,转换可能会导致在内存中同时保留两个版本的数据,这在某些操作场景下可能会成为问题。不过不用担心!接下来我们将介绍一些策略来减轻这个问题。

防止 pandas 占用过多资源

在前一节中,我们看到从 CSV 数据的 Arrow 表创建 pandas DataFrame 需要超过 500 毫秒。如果你觉得这个速度有些慢,那是因为它必须复制我们数据中的所有字符串。用于将 Arrow 表和数组转换为 DataFrame 的函数提供了一个参数 zero_copy_only,如果将其设置为 true,当转换需要复制数据时,将会抛出一个 ArrowException 错误。这种全有或全无的策略仅适合在你需要严格控制内存使用的情况下使用。要实现零拷贝转换,必须满足以下条件:

  • 整数(有符号或无符号)数据类型,无论位宽,或者是浮点数据类型(float16、float32 或 float64)。这也涵盖了使用这些数据表示的各种数值类型,如时间戳、日期等。
  • Arrow 数组中没有 null 值。Arrow 数据使用位图来表示 null 值,而 pandas 不支持这种表示方式。
  • 如果是分块数组,那么必须只有一个块,因为 pandas 要求数据完全连续。

PyArrow 库提供了两个选项来限制转换期间潜在的数据拷贝问题——split_blocksself_destruct。由于 pandas 底层使用 NumPy 进行计算,它喜欢将相同数据类型的列收集到二维 NumPy 数组中,因为这可以加速对多列进行的批量操作,如同时计算多个列的总和。下图简要展示了 pandas 中 DataFrame 的内存管理方式。pandas 中有一个称为 Block Manager 的对象,负责内存分配并跟踪底层数据数组的位置。不幸的是,如果你逐渐按列构建 DataFrame,Block Manager 会将这些单独的列合并到称为块(blocks)的组中,而这种合并会要求在内部复制数据以便创建该块:

image.png

PyArrow 库非常努力地构建 pandas 期望的整合块,以避免在转换为 DataFrame 后执行额外的分配或复制。然而,这样做的缺点是需要从 Arrow 复制数据,这意味着数据的内存使用峰值可能会达到数据总大小的两倍。前面提到的 split_blocks 选项会为每一列生成一个单独的块,而不是提前进行整合。如果设置为 True,它将加快转换过程,并可能避免最糟糕的情况(即完全双倍的内存使用)。设置此选项后,如果数据符合零拷贝转换的条件,你将获得真正的零拷贝操作。

让我们看看这个选项的实际效果: 首先,我们需要导入所需的库:pandas、PyArrow 和 NumPy:

import pandas as pd
import pyarrow as pa
import numpy as np

然后,我们创建一组随机的浮点数据作为 NumPy 数组来进行测试:

nrows = 1_000_000
ncols = 100
arr = np.random.randn(nrows)
data = {'f{}'.format(i): arr for i in range(ncols)}

现在,我们对转换进行计时,看看结果如何:

In [8]: %timeit df = pd.DataFrame(data)
157 ms ± 13.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [9]: %timeit df = pa.table(data).to_pandas()
115 ms ± 4.91 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [10]: %timeit df = pa.table(data).to_pandas(split_blocks=True)
3.18 ms ± 37.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

太神奇了!当然,如果你在这些转换之后执行了一堆 pandas 操作,那么节省的 100 毫秒可能会在接下来的整合过程中重新出现,但这个数字仍然非常令人印象深刻!即使在 Arrow 执行整合的情况下,转换的速度仍比直接从 NumPy 数组创建 DataFrame 快了大约 26%。Arrow 进行这些优化的原因之一是为了防止大家都需要为 DataFrame 编写转换器。各种组件、工具和系统可以生成任意语言的 Arrow 格式数据(即使它们不直接依赖于 Arrow 库),然后使用 PyArrow 将其转换为 pandas DataFrame。你无需编写转换器,因为 Arrow 库的速度可能比任何自定义转换代码都快,而且会减少维护代码的负担,真是双赢!

那么听起来有点危险的 self_destruct 选项呢?通常,当你复制数据时,内存中会存在两份副本,直到变量超出作用域并且 Python 的垃圾回收机制清理它们。使用 self_destruct 选项将逐列销毁内部的 Arrow 缓冲区,这样每转换一列就会将内存释放回操作系统。关键要记住的是,转换后你的 Table 对象将不再安全使用,调用它的方法会导致 Python 进程崩溃。你还可以同时使用两个选项,在某些情况下,这将显著减少内存使用:

>>> df = table.to_pandas(split_blocks=True, self_destruct=True)
>>> del table # 不必要的步骤,但作为好习惯可以加上

注意: 使用 self_destruct 不一定能节省内存!由于转换是在每列转换时进行的,内存释放也是逐列发生的。在 Arrow 库中,多个列很可能共享同一个底层内存缓冲区。在这种情况下,直到引用同一缓冲区的所有列都转换完毕,内存才会被释放。

尽管 pandas 在数据科学领域使用广泛,但它也有一些缺点。这些缺点主要集中在内存使用、处理 null 数据以及多核执行方面。随着时间推移,其他 DataFrame 库不断涌现以弥补这些空缺。既然我提到了,你可能已经猜到接下来我们要介绍什么了:那就是 Polars 库。

极地熊用生锈的箭

与 pandas 一样,Polars 库的基本构建块是 DataFrame。我们从最简单的情况开始:一个单列,并将 Polars 的 DataFrame 转换为 Arrow 表,再从 Arrow 表转换回来:

>>> import pyarrow as pa
>>> import polars as pl
>>> df = pl.DataFrame({"a": [1, 2, 3]})
>>> table = df.to_arrow()  # 转换为 Arrow
>>> df_new = pl.from_arrow(table)  # 转换回 Polars

除了 CategoricalType 以外,Polars 中的列都是零拷贝地转换为 Arrow 的,基本上它只是重新分配指针并复用同一块内存! 与 pandas 一样,Polars 也有与 Arrow 数据类型对应的数据类型,不仅包括基本类型,还包括嵌套类型,如 StructList。不过,这里有一个重要的注意事项:对于使用偏移量的类型(如 ListStringBinary),Polars 总是使用大偏移量(即 64 位偏移量)变体。这意味着从 32 位偏移量变体转换时,将会执行复制并将偏移量转换为 64 位值。

如前所述,Polars 是用 Rust 语言实现的,主要追求性能,类似于 PyArrow 使用 C++ Arrow 库实现。因此,它也有用于将文件读取为 DataFrame 的 API。让我们看看它与之前用 pandas 和 PyArrow 做的比较如何:

In [1]: import pyarrow as pa
In [2]: import polars as pl
In [3]: import pyarrow.csv
In [4]: %timeit table = pa.csv.read_csv('train.csv')
90.5 ms ± 1.18 ms 每次循环 (平均 ± 标准偏差,7 次运行,每次 10 次循环)
In [5]: %timeit df = pl.read_csv('train.csv')
128 ms ± 1.01 ms 每次循环 (平均 ± 标准偏差,7 次运行,每次 10 次循环)
In [6]: df = pl.read_csv('train.csv')
In [7]: table = pa.csv.read_csv('train.csv')
In [8]: %timeit table = df.to_arrow()
42.9 µs ± 218 ns 每次循环 (平均 ± 标准偏差,7 次运行,每次 10,000 次循环)
In [9]: %timeit df = pl.from_arrow(table)
108 ms ± 1.4 ms 每次循环 (平均 ± 标准偏差,7 次运行,每次 10 次循环)

表现不错吧?

之前我们看到,pandas 读取相同 CSV 文件的耗时大约为 3.49 秒。而直接用 PyArrow 读取则只需 90 毫秒。使用 Polars 读取文件到 DataFrame 的耗时略高,为 128 毫秒,而将其转换回 Arrow 表仅需 42 微秒;或者将 Arrow 表转换回 Polars DataFrame 则只需 108 毫秒。总体来看,速度非常快!根据你的机器和 Polars 版本,可能会稍快或稍慢于 PyArrow,但总体来说,它们的性能大致相当。你可以查看文档,找到适用于其他文件类型(如 Parquet 和 JSON)的等效方法 (docs.pola.rs/user-guide/io/)。

到目前为止,我们已经多次提到零拷贝,并介绍了 Arrow 如何通过各种方式促进数据的高效传输。Arrow 的列式格式使其非常容易流式处理和传输内存中的原始缓冲区,或者重新利用它们以提高性能。通常,传递数据时需要进行序列化和反序列化,但正如前面提到的,Arrow 允许你跳过这些序列化和反序列化的开销。

为了在处理这些流和文件时确保高效的内存使用,Arrow 库提供了各种内存管理的实用工具,并在其内部被广泛使用。接下来,我们将介绍这些辅助类和工具,帮助你利用它们让你的脚本和程序更加简洁高效。我们还将讨论如何在编程语言之间共享缓冲区以提高性能,以及为什么你可能想这么做。

分享就是关怀……尤其是你的内存

前面我们提到过对 Arrow 数组进行切片的概念,它允许你在不复制数据本身的情况下,获取表、记录批次或数组的视图。这种操作甚至可以深入到 Arrow 库使用的底层缓冲区对象,这些缓冲区对象可以被消费者使用,即使他们不直接处理 Arrow 数据,也能有效管理他们的内存。Arrow 库通常提供内存池对象来控制内存的分配,并跟踪 Arrow 库分配了多少内存。这些内存池随后被数据缓冲区和 Arrow 库中的其他组件使用。

深入内存管理

继续我们的 Go、Python 和 C++ Arrow 实现的研究,它们都提供了类似的内存池来管理和跟踪内存使用情况。以下是一个内存池的简化示意图:

image.png

随着内存需求的增加,内存池会随着更多内存的分配而扩展。当内存被释放时,它将回收到池中,以便在未来的分配中重新使用。具体的管理策略会根据不同的实现而有所不同,但基本理念如前面的图示所示。内存池通常用于生命周期较长且占用内存较大的数据,比如数组和表的数据缓冲区,而小的临时对象和工作区将使用编程语言中的常规分配器。

在大多数情况下,默认的内存池或分配器会被使用(如前面多个代码示例中所见),但许多 API 允许你传入一个特定的内存池实例来执行分配操作,接下来将对此进行介绍。

C++ 实现

库提供了 arrow::MemoryPool 类,用于操作或检查内存的分配。在库首次初始化时,进程范围内的默认内存池会被初始化。可以通过代码中的 arrow::default_memory_pool 函数访问该内存池。根据库的编译方式以及 ARROW_DEFAULT_MEMORY_POOL 环境变量,默认池可能由 jemalloc 库、mimalloc 库或标准的 C 语言 malloc 函数实现。内存池本身提供了一些函数,可以手动将未使用的数据尽力释放回操作系统(仅在底层分配器保留未使用内存的情况下),报告内存池的峰值分配量,以及返回当前分配但尚未通过池释放的字节数。

内存分配器

使用 jemallocmimalloc 等自定义分配器的好处在于,它们可能显著提升性能。根据基准测试结果,两者都显示出比传统的 malloc 更低的系统内存使用和更快的分配速度。值得对你的工作负载进行不同分配器的测试,看看是否能从中受益!

数据缓冲区的操作

对于操作数据缓冲区,提供了 arrow::Buffer 类。类似于标准容器(如 std::vector),可以通过 ResizeReserve 方法预先分配缓冲区,使用 BufferBuilder 对象来操作。这些缓冲区将被标记为可变或不可变,具体取决于它们的构造方式,指示它们是否可以调整大小和/或重新分配。如果你使用输入流(InputStream)等 I/O 功能,建议使用提供的 Read 函数将数据读入 Buffer 实例,因为在许多情况下,它可以对内部缓冲区进行切片,而无需复制额外的数据。下图展示了一个已分配的缓冲区,其中包含长度和容量,以及缓冲区的一个切片视图。该切片知道它并不拥有指向的内存,因此在清理时不会尝试释放该内存:

image.png

Python

由于 Python 库是基于 C++ 库构建的,因此前面提到的关于内存池和缓冲区的所有功能也在 Python 库中可用。pyarrow.Buffer 对象封装了 C++ 的缓冲区类型,以允许其他更高层次的类与它们可能拥有或不拥有的内存进行交互。缓冲区可以通过切片和内存视图相互引用,从而与其他缓冲区建立父子关系,这样内存可以在不同的数组、表和记录批次之间轻松共享,而无需复制数据。在需要 Python 缓冲区或内存视图的地方,可以使用 Buffer 来进行零拷贝数据操作,而不必复制数据:

>>> import pyarrow as pa
>>> data = b'helloworld'
>>> buf = pa.py_buffer(data)
>>> buf
<pyarrow.lib.Buffer object at 0x000001CB922CA1B0>

调用 py_buffer 函数时不会分配新的内存。它只是对 Python 已为 bytes 对象分配的内存进行零拷贝视图。如果需要 Python 缓冲区或内存视图,则可以通过缓冲区进行零拷贝转换:

>>> memoryview(buf)
<memory at 0x000001CBA8FECE88>

最后,缓冲区上有一个 to_pybytes 方法,可以创建一个新的 Python 字节串对象。这将复制缓冲区引用的数据,确保新的 Python 对象与缓冲区之间的内存隔离。

同样,由于一切都由 C++ 库支持,Python 库也有自己的默认内存池,可以告诉你目前分配了多少内存。我们可以分配自己的缓冲区并观察这种情况:

>>> pa.total_allocated_bytes()
0
>>> buf = pa.allocate_buffer(1024, resizable=True)
>>> pa.total_allocated_bytes()
1024
>>> buf.resize(2048)
>>> pa.total_allocated_bytes()
2048
>>> buf = None
>>> pa.total_allocated_bytes()
0

你可以看到,一旦内存被垃圾回收,内存就会被释放,内存池会反映出该内存已不再分配。

重要提示:时机问题

垃圾回收和内存释放不是即时的,因此对于较大的内存分配,可能需要使用 sleep 函数来稍微延迟,以便解释器和操作系统有时间处理并同步内存释放。

Golang

与 Python 和 C++ 库类似,Go 库也通过 memory 包提供了缓冲区和内存分配管理。memory.DefaultAllocator 是默认分配器的实例,即 memory.GoAllocator。由于分配器定义是一个接口,因此可以根据项目的需要轻松构建自定义分配器。如果可用 C++ 库,则可以在使用 Go Arrow 库构建项目时提供 "ccalloc" 构建标签。在这种情况下,可以使用 CGO 提供 NewCgoArrowAllocator 函数,该函数使用 C++ 内存池对象而不是 Go 默认分配器来分配内存。如果需要在 Go 和其他语言之间传递内存,这一点非常重要,因为它可以确保 Go 的垃圾收集器不会干扰内存管理。

最后,memory.Buffer 类型是 Go 库中的主要内存管理单元。它的工作方式与 C++ 和 Python 库中的缓冲区类似,提供对底层字节的访问,能够调整大小,并在包装字节切片时检查其长度和容量。

性能优化的缓冲区管理

通过这种内存和缓冲区管理,我们可以想象几个场景,将这些功能结合起来可以确保卓越的性能,如下所示:

假设你想对包含数十亿行的非常大的数据集进行分析。提高此类操作性能的一种常见方法是对行的子集进行并行操作。通过能够在不复制底层数据的情况下对数组和数据缓冲区进行切片,这种并行化操作变得更快,并且内存需求更低。你操作的每个批次并不是数据的副本,而只是数据的一个视图,如下图所示。切片列上的虚线表示它们只是各自列中数据子集的视图,正如之前的切片缓冲区图所演示的那样。每个切片都可以安全地并行操作:

image.png

  • 也许你有一系列的数据列,并希望逐步过滤掉那些每列都是空值的行。最简单的方法是遍历每一行,如果在某个索引处至少有一个非空值,则将该行的数据复制到每列的新版本中。如果你正在处理嵌套列,这可能会变得更加复杂。然而,当使用 Arrow 数组时,你可以通过使用有效性位图缓冲区来加快这一过程!只需对所有位图执行按位或操作,就可以得到一个单一的位图,它代表了最终的过滤索引。然后,你不必为每一列逐步构建一个过滤后的副本,而是可以逐列进行处理,以获得更好的 CPU 缓存命中率和内存局部性。下图直观地展示了这个过程。根据数据的总大小和结果中的行数,可能更有意义的是仅对每组行进行切片,而不是将它们复制到新的列中。无论哪种方式,你都可以控制内存的使用和释放时机:

image.png

如果你能够传递某个数据缓冲区的地址,并且知道 Arrow 的内存格式与编程语言无关,那么只需要一些元数据,你就可以在不同的运行时和编程语言之间共享数据表。你可能会问,为什么要这样做呢?让我们来看一下这样做有多大的用处……

跨越边界

以下图表展示了数据科学中常见的工作流程之一:

image.png

该工作流程的步骤如下:

  1. 从一个或多个结构化查询语言 (SQL) 引擎中,通过某个库使用 Java 数据库连接 (JDBC) 驱动从 Python 进程中查询数据。
  2. 查询到的数据现在存储在 JVM 内存中,接着被复制到 Python 内存中。
  3. 现在,数据已存储在 Python 内存中,随后使用 pandas 进行某种处理。
  4. pandas 处理的分析结果被馈送到所使用的模型中,如机器学习模型或其他统计模型。

对于大数据集而言,此工作流程中最昂贵的部分是将数据从 JVM 复制到 Python 内存,并在 pandas 中将数据从行转化为列的格式。为了改进此类工作流程,Arrow 库提供了一个稳定的 C 数据接口,允许在不复制数据的情况下跨越这些边界共享数据,直接通过内存指针进行共享。通过这种方式,数据指向的是内存,而无需创建大量中间的 Python 对象。这个接口由几个简单的头文件定义,可以被复制到任何能够与 C API 通信的项目中,比如通过外部函数接口 (FFI) 来实现。

在此工作流程中,Java 库中的 JDBC 适配器用于检索查询结果,将行转换为 JVM 中的列,并将数据作为 Arrow 记录批次存储在 JVM 自身不管理的堆外内存中。然后,此原生内存布局可以使用 C 数据接口向 PyArrow 库提供指向原始数据缓冲区和逻辑结构的指针,使库能够正确解释并使用内存。下图展示了使用这些接口的新工作流程:

image.png

这一次,工作流程如下:

  1. 从一个或多个 SQL 引擎中,通过某个库使用 JDBC 驱动从 Python 进程中查询数据。
  2. 当数据行返回时,它们被转换为列格式的 Arrow 记录批次,并且原始数据存储在 JVM 管理之外的“堆外”内存中。
  3. 数据无需复制,Arrow 数据可以通过 JVM 已分配的内存指针直接在 Python 中访问。不会创建中间的 Python 对象,也不会对原始数据进行任何复制。
  4. 如前所述,我们可以将 Arrow 记录批次转换为 pandas DataFrame。
  5. 最后,将这些 DataFrame 传递给所需的库,用分析过的数据填充模型。

虽然这看起来变化不大,但在实践中,这可以带来巨大的性能提升。以 Dremio 作为 SQL 引擎,并使用样例 NYC Taxi Trip Duration 数据集,我对这两种方法的性能进行了比较:

  • 采用传统方法,从 Dremio 中查询约 62.3 万行和 6 列的数据集,并创建 DataFrame,平均耗时 1 分 5 秒。
  • 采用共享内存避免数据复制的方法,同样的查询耗时约 573 毫秒。这比传统方法快了约 113 倍,提升了 11,243.8%。

如果你的结果集足够小,使用共享内存方法的好处不会特别明显,可能不值得为此增加复杂性和依赖性。下图展示了这两种方法在不同行数条件下的性能表现,每次查询均包含 6 列数据。我们可以看到,如果行数少于 10,000 行,即使相对数字显示出显著加速,但绝对耗时并不多,这要根据具体的工作流程来判断其是否有益。

image.png

出于纯粹的好奇心,我测试了这两种工作流程在我使用的完整 Parquet 数据集上的表现,这个数据集有超过 6200 万行数据。使用传统的复制方法,整个过程花费了超过 3 小时,而利用跨 C 数据接口的共享内存工具仅用了约 58.7 秒。这个结果令人惊叹,性能提升约 184 倍,改进幅度约 18,520%!

重要提示 需要注意的是,稍后在第 8 章中,我们会讨论一种 Arrow 原生的 JDBC 替代方案,名为 Arrow 数据库连接(ADBC)。ADBC 利用我们在此提到的 C 数据接口(在第 4 章《通过 Arrow C 数据 API 跨越语言障碍》中详细介绍)来提供极高性能的数据库交互接口。

如果你还没有猜到,C 数据接口的主要目标用户是那些构建使用 Arrow 的库、工具和实用程序的开发人员。已经有多个软件包利用了这些接口,比如 arrow R 包中的 reticulate 方法(rstudio.github.io/reticulate/… R 和 Python 之间的数据,以及我之前使用过的 pyarrow.jvm 模块。随着越来越多的开发者和库构建者利用 C 数据接口来通过共享内存传递数据,我们将看到常见数据任务的整体性能大幅提升,把更多的 CPU 周期和内存留给真正需要的分析计算,而不是反复复制数据以便在你想使用的工具中访问它。

如果你是那些库和工具开发者之一,或者是处理数据传递的工程师,不妨利用并尝试一下这个接口。除了处理原始数据外,C 接口还支持流式数据,这样你可以将记录批次直接流式传输到共享内存中,而不是进行复制。截至本文撰写时,已经可以在 Arrow 的 C++、Python、R、Rust、Go、Java、C/GLib 和 Ruby 实现中使用 C 数据接口。快去利用这个在工具之间共享数据的超强方法吧!行动起来!

针对库开发者的说明 为了方便库生产和消费 Arrow 数据,一个非常小巧的 C 库被创建,旨在为那些想要使用 Arrow 但又不想链接到完整 Arrow C++ 库的项目服务。如果你对此感兴趣,可以查看这个既小巧又非常实用的库:github.com/apache/arro…

总结

到此为止,你应该已经相当熟悉了 Apache Arrow 库的各种主题和概念,并且知道如何将它们集成到日常的工作流中。不论你是利用文件系统抽象、数据格式转换,还是零拷贝通信的优势,Arrow 都可以融入任何数据工作流的多个部分。在继续学习之前,请确保你理解了迄今为止涉及到的格式、通信方法以及 Arrow 库提供的实用工具的概念。通过多次实践,尝试使用不同的策略来管理你的数据,并在工具和实用程序之间传递数据。如果你是一名构建分布式系统的工程师,尝试使用 Arrow 的进程间通信(IPC)格式(下一章将详细介绍!),并与之前的数据传递方式进行对比。哪个更易用?哪个性能更高?

下一章《格式和内存处理》将深入探讨数据传递和存储的各种方式,讨论它们之间的关系和应用场景。目标是了解 Arrow 及其 IPC 格式如何适应现有的数据生态系统,并分析各种实现方式的权衡利弊。准备好了吗?那我们开始吧!