使用 Apache Arrow 进行内存分析——在机器学习工作流中使用 Arrow我们刚刚讨论了如何使用 Arrow 数

我们刚刚讨论了如何使用 Arrow 数据库连接 (ADBC)，它提供了一种与多种数据源进行高效交互的方式。在本章中，我们将探讨如何使用这些数据：即机器学习（ML）。ML 绝不仅仅是一个流行词汇——它经常被用于模式识别、数据驱动的决策制定以及生成式人工智能（GenAI）系统。虽然这可能是个有争议的观点，但从本质上讲，ML 工作流只是标准数据管道的一种特化形式。因此，凡是涉及数据处理的地方，Arrow 都有机会发挥极大的作用！

无论您是在进行特征工程、模型训练、预处理还是其他任务，许多常用工具和实用程序都提供了与 Arrow 的互操作性。其中一些工具甚至在底层使用了 Arrow。

在本章中，我们将涵盖以下主题：

使用 Arrow 和一些配置设置来提升 Spark 数据管道的性能
Hugging Face 如何使用 Arrow 处理超出内存的数据集
如何使用 PyTorch 和 TensorFlow 等标准工具与 Arrow 数据集进行协作
利用标准实现 ML 库与 GPU 的互操作性

技术要求

在本章中，我们将包含大量的技术示例。因此，您需要以下条件：

一台联网的计算机。
一个运行 Jupyter 的单节点或多节点 Apache Spark 集群。Docker 是最简单的设置方式，我们将在本章中介绍如何完成设置。
安装了 pyarrow 模块及其 dataset 子模块的 Python 3.8+ 版本。
您喜欢的代码编辑器，如 Emacs、Vim、Sublime 或 VS Code。

在 Jupyter 中激发新的想法

Apache Spark 是一个开源的分析引擎，专为在大型集群上进行分布式处理而设计，利用这种架构带来的并行性和容错能力。我认为它是自 JavaScript 发明以来，既被深深喜爱又被强烈批评的软件之一！对它的喜爱来自于其所启用的工作流，但它也因其易碎性和使用难度而臭名昭著。如果你对 Spark 不熟悉，它通常与 Scala、Java、Python 和/或 R 结合使用，并且能够运行分布式 SQL 查询。由于 Python 上手简单，编写速度快，数据科学家常常使用 Jupyter Notebook 结合 Python 快速创建和测试分析模型。

这种工作流风格非常适合快速迭代各种想法，并验证其可行性和正确性。然而，工程师和数据科学家常常会被限制在这样一个现实中：Python 的运行速度相对较慢。此外，除非你有一个拥有大量核心的大型集群，否则根据具体的计算和使用场景，Spark 也可能运行得相当慢。这就引出了一个问题：我们如何能够在提高编写计算代码的简便性和提高这些计算性能之间取得平衡？答案显然是利用 Arrow 和 Parquet 的集成，借助列式格式来提升性能。

理解 Spark 中的 Arrow 集成

首先，我要指出，这并不是 Spark 和 Arrow 的对立关系，而是如何利用 Arrow 来增强现有的 Spark 管道。许多 Apache Spark 管道不会从使用 Arrow 中受益，因为 Spark 具有与 Arrow 不同的内存 DataFrame 格式，在两者之间进行转换会导致性能下降，因此所有的优势都需要谨慎权衡。然而，这种结合在某些场景下运作得非常完美，例如在你的 Spark 管道中使用 pandas 自定义函数时。在这种情况下，Spark 可以利用 Arrow 来转换和传递数据，从而享受本书之前介绍的各种好处。图 9.1 显示了一个简化的 PySpark 工作流程图，以说明我们讨论的内容：

当你运行 PySpark 时，会为你启动两个进程：Python 解释器本身和一个 JVM 进程。这是因为 Spark 是用 Scala/Java 编写的，而不是 Python。Spark 所做的所有繁重的计算工作都是在 Python 进程之外完成的，Python 通过 Py4J 桥接提供了一个接口，用来向 Spark 发送命令。问题在于，当你想在 Python 中与数据交互，然后将其发送到 Spark，再获取结果时，必须通过某种方式将数据传递过这个桥接。想象一下，在 Python 中加载一个 4 GB 的 pandas DataFrame，稍作处理后将其发送到 Spark 进行计算和进一步分析。图 9.2 展示了这个过程在底层发生的情况：

在上图中，你可以看到 4 GB 的数据需要被序列化为字节流，然后传递给 JVM 进程，接着反序列化，以便 Spark 进行操作。而 Spark 的计算结果则要经过相反的过程。当你处理大数据集时，优化这个数据传输过程可以为你节省大量的时间和计算资源！同时要注意，由于 pandas DataFrame 和 Spark DataFrame 在内存中的表示方式不同，因此在这两者之间仍然需要进行转换。而这正是 Arrow 能加速该过程的地方。

本质上，每当你想使用一个不是 Spark 原生内存格式的库时，你都需要在这些格式之间进行转换。这包括使用一些 Java 库以及任何非 Java/Scala 库（例如运行 Python 自定义函数）。有些操作在使用 Arrow 格式时比 Spark 格式更快，反之亦然，但在大多数情况下，只有当你在非 Spark 格式（如 pandas）中进行大量工作时才值得使用 Arrow，而数据科学家经常这样做，因为 pandas 比 Spark 更友好且易于使用。

在这个示例中，我们将使用之前使用过的免费和开源的 NYC Taxi 数据集中的一个文件片段。你可以在本书的 GitHub 仓库的 sample_data 文件夹中找到这些数据：GitHub 链接。我故意只获取了一个文件的片段而不是整个文件，因为我们仅展示一些容易运行并能在本地使用 Docker 复制的示例，而不是要求你使用自己的 Spark 集群。我们将研究两个用例：

将数据从原始 CSV 文件导入 Spark 并导出数据
对数据集中的一个或多个数值字段执行归一化计算

在深入代码之前，让我们先使用 Docker 启动开发环境。

容器化让生活更轻松

在第7章《探索 Apache Arrow Flight RPC》中，我们已经介绍了如何安装和使用 Docker。如果你需要回顾一下，请返回该章节的“每个人都有容器化的开发环境！”部分。一旦 Docker 安装完毕，你可以使用以下命令启动我们将要使用的开发容器：

$ docker run -it --name jupyter -v ${PATH_TO_SAMPLE_DATA}:/home/jovyan/work -e JUPYTER_ENABLE_LAB=yes -p 8888:8888 -p 4040:4040 jupyter/pyspark-notebook

这里，PATH_TO_SAMPLE_DATA 应该是一个环境变量，包含本书 GitHub 仓库本地克隆的路径，该仓库中包含名为 chapter9 的目录内的 Jupyter Notebook。此命令将启动 Docker 镜像，并将其日志直接输出到你的终端，同时绑定本地端口 8888 和 4040 供使用。如果你愿意，可以选择不同的本地端口进行绑定。

启动后，确保查看日志，因为你需要从中获取启动 Jupyter 会话的 URL。图9.3 显示了你需要在日志中查找的内容，突出显示的行是你需要复制并粘贴到浏览器中的 URL，以访问 Jupyter 并打开仓库中的 notebook。

打开 chapter9.ipynb 后，你会看到 notebook 的第一个单元格，它为你设置 PySpark 环境。我们使用的 Docker 镜像已经包含了 pyarrow 模块，因此可以直接使用。点击第一个单元格并按下 Shift + Enter，你将启动 Spark 的主节点（master）和执行器（executor），同时下载从 Spark 访问 AWS S3 所需的包。运行结果应类似于图 9.4 所示：

此时，你已经准备好运行我将带你体验的示例了。让我们来看一下如何从 Spark 中集成 Arrow 获益。

通过 Arrow 和 PySpark 点燃快乐

在提供的文件中，有两个文件值得关注：sliced.csv 和 sliced.parquet。它们是我们将在这些示例中使用的 NYC Taxi 数据集的片段。

设置数据

sliced.parquet 文件位于本书的 GitHub 仓库的 sample_data 目录中。第9章的 Jupyter Notebook 的第一个单元格包含一些快速代码，用于从 Parquet 文件中写出 CSV 文件。这样，你无需下载一个大文件。

CSV 文件大约有 511 MB，而 Parquet 文件只有 78 MB。值得注意的是，它们包含相同的数据！这种差异是我们在之前章节中已经见过的：这就是二进制列式存储格式以及它所利用的压缩的优势所在。除了展示如何从 CSV 文件将数据导入 Spark，我们还将展示从 Parquet 文件获取相同数据的速度要快得多。

步骤 1 – 使数据可用

首先，我希望你回顾第2章《使用关键的 Arrow 规范》中的代码示例，并尝试将 sliced.csv 文件读取到 pandas DataFrame 中。记住，按 Shift + Enter 可以运行单元格中的代码，不过如果你更喜欢直接使用交互式控制台，也可以切换过去。像之前一样，你可以在执行之前在代码行前加上 %time 来获取执行后的计时信息。对于多行单元格，可以使用两个百分号 %%time。为了防止你忘记，这里提供了代码，代码可以在 Jupyter Notebook 中找到：

%%time
import pyarrow as pa
import pyarrow.csv
pdf = pa.csv.read_csv(
    '../sample_data/sliced.csv').to_pandas()

当我运行这段代码时，得到了以下输出。根据你运行的机器配置，具体的计时结果可能会有所不同：

CPU times: user 4.11 s, sys: 1.67 s, total: 5.78 s
Wall time: 383 ms

注意：
如果你对这些术语不熟悉，wall time 是用秒表直接计时的整个过程的耗时。user time 是 CPU 实际使用的时间；在多核机器上，这个时间可能比 wall time 大得多。在本例中，我们知道 PyArrow 默认会使用多线程来读取 CSV 文件，因此 user time 比 wall time 大得多。sys time 则是内核执行系统级操作（如上下文切换、磁盘 I/O 和资源分配）所耗费的时间。

你可能会注意到，这与我之前的计时输出有所不同，差异在于使用了 %%timeit 而不是 %%time。%%timeit 会多次运行命令，给出运行时间的平均值和标准差，并告诉你它运行了多少次。使用 %%timeit 运行同样的代码，我得到了以下输出：

vbnet
复制代码
385 ms ± 33.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

无论如何，我们可以看到，使用 PyArrow 将这个 CSV 文件读入 pandas DataFrame 平均只需不到半秒的时间（在我的笔记本电脑上）。但请记住，我们需要一个 Spark DataFrame 来供 Spark 使用！要实现这一点，可以在 SparkSession 对象上使用一个非常有用且名称贴切的方法，称为 createDataFrame。然而，对于一个包含 350 万行和 21 列的 pandas DataFrame，这可能是一个非常昂贵的操作。在我的机器上，这个过程花费了超过一个小时！那么如果我们直接使用 Spark 自身的函数来读取呢？让我们在 Jupyter Notebook 的某个单元格中试试这个方法：

%%time
df = spark.read.format('csv').load(
    '../sample_data/sliced.csv', inferSchema='true', header='true')

我在机器上得到的输出如下：

CPU times: user 346 µs, sys: 4.22 ms, total: 4.56 ms
Wall time: 2.5 s

相当大的差别，对吧？这就是我之前提到的不同 DataFrame 格式之间转换的成本。话虽如此，Spark 支持利用 PyArrow 来增强这些操作并减少数据复制。你可以通过运行以下命令启用这种行为：

spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

就是这样。有了这个配置设置后，我们可以看到执行时间有了巨大的变化。运行以下代码：

%%time
df = spark.createDataFrame(pa.csv.read_csv(
    '../sample_data/sliced.csv').to_pandas())

输出如下：

CPU times: user 5.49 s, sys: 3.39 s, total: 8.87 s
Wall time: 2.23 s

通过启用 Spark 中的 Arrow 优化，我们可以利用 Arrow 以极快的速度读取 CSV 文件（比在 Spark 中通过单个执行器原生读取要快），但最终仍然能够得到一个 Spark DataFrame。结果是，我们可以在大约 2 秒内获取到一个 DataFrame，而不是花费超过一个小时。

虽然这是将 pandas DataFrame 转换为 Spark DataFrame 的一个相对简单的示例，但随着越来越多直接与 Arrow 格式兼容的库和模块的出现，你的预分析操作（例如数据清理或数据处理）通常可以让你最终获得一个易于兼容的 DataFrame。与 Spark 的集成确保了在某些情况下获得更高的性能，这些情况下由于数据管道配置的原因，直接将数据原生地读入 Spark 可能不可行。

请记住
这仍然是一个仅在我本地机器上使用单个执行器的示例。Spark 的优势在于可以跨多个机器和核心进行并行化，而在这里我没有使用。你仍然可以通过使用 Arrow 库来受益，但它并不那么直接。对于数据的 Parquet 版本也是如此。如果你的数据集足够大，应该使用单独的 Spark 任务来创建多个较小的 DataFrame，并使用 block_size 和 skip_rows_after_names 选项来读取 CSV，或者从 Parquet 文件中读取较小的行组。

由于优化和延迟加载的原因，仅通过读取 Parquet 文件，速度提升并不那么明显。无论是直接使用 PyArrow 还是使用 Spark，它都会在不到一秒的时间内将 Parquet 文件读取到 DataFrame 中。不过，如果我们使用 describe().show() 来强制 Spark 读取整个文件并对其进行一些操作，我们可以看到使用 PyArrow 读取 Parquet 文件比 Spark 原生读取器更快的好处：

%%time
df = spark.read.format('parquet').load('../sample_data/sliced.parquet')
df.describe().show()

我得到以下输出，省略了具体的数据，只展示报告的 CPU 和 wall time：

...
CPU times: user 6.01 ms, sys: 997 µs, total: 7.01 ms
Wall time: 16.5 s

从 Spark 原生读取器的计时结果中，我们可以看到工作主要由 Spark 的 driver 和 executors 完成，而不是 Python 进程。尽管 wall time 为 16.5 秒，但报告的 CPU 时间只有几毫秒，因为所有工作都是由 JVM 进程和 Spark 进程完成的，结果从 JVM 返回到 Python 进程。如果我们使用 PyArrow 来读取文件，则会看到差异：

%%time
df = spark.createDataFrame(pq.read_table('../sample_data/sliced.parquet').to_pandas())
df.describe().show()

再次省略原始数据，得到以下计时输出：

...
CPU times: 1.86 s, sys: 1.08 s, total: 2.94 s
Wall time: 7.61 s

我们的计时显示 Python 进程花费了总共 2.94 秒，而之前是 7.01 毫秒。这是因为读取 Parquet 文件的工作是由 PyArrow 模块在 Python 中完成的，但由于启用了 Spark 中的 Arrow 和零拷贝优势，数据转换过程非常快。如果我们将两者的 Spark 执行计划的可视化进行对比，如图 9.5 所示，我们可以理解发生了什么：

对比执行计划时，我们可以看到不同之处在于 ParallelCollectionRDD 和 Arrow toDataFrame 步骤，而不是 FileScanRDD。看起来，当流式传输 Arrow 记录批次时，Spark 在并行化任务方面做得更好，而不是直接进行读取。pyarrow 模块将整个 Parquet 文件读取到内存中，然后使用 Arrow 的进程间通信（IPC）格式将数据传递给 Spark，以便于并行化处理。如果我们查看这两种情况下的执行时间线，这一点就更加明显，如图 9.6 所示：

你能猜出哪个截图是使用 Arrow 读取文件并传递 DataFrame 的运行结果吗？我给你几秒钟……是的，是上面的截图。你看到的是 Spark 如何选择将任务分配到不同的工作进程中，驱动程序使用了这些进程。截图中的颜色代表了执行器如何消耗时间。蓝色部分是调度延迟，红色部分是任务反序列化的时间。而我们最感兴趣的是绿色部分——它表示执行任务本身所花费的时间。

在这两种情况下，Spark 都将工作分成了 12 个任务，但在流式传输已经读取的 Arrow 记录批次的情况下，而不是直接读取文件，工作分布得更加均匀，每个执行器都承担了一部分工作。而当 Spark 自己读取文件时，它只是并行读取了 Parquet 文件，将所有数据收集到一个执行器中进行所有计算，导致了额外的计算时间。

接下来，我们只需要一部分列，这进一步提高了我们的读取性能：

首先是 Spark 原生版本：

%%timeit
df = spark.read.format('csv').load(
    '../sample_data/sliced.csv',
    inferSchema='true',
    header='true').select('VendorID',
        'tpep_pickup_datetime', 'passenger_count',
        'tip_amount', 'fare_amount',
        'total_amount')

输出结果：

1.6 ± 66 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

然后，使用 PyArrow，我们可以将数据转换为 pandas DataFrame：

%%timeit
df = spark.createDataFrame(pa.csv.read_csv(
    '../sample_data/sliced.csv',
    convert_options=pa.csv.ConvertOptions(
        include_columns=['VendorID',
            'tpep_pickup_datetime',
            'passenger_count', 'tip_amount',
            'fare_amount', 'total_amount'])
    ).to_pandas())

输出结果：

577 ms ± 1.43 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

通过利用 Arrow 库更好的读取性能和零拷贝转换，我们可以展示出与使用原生 Spark 加载函数相比，性能显著提升。

步骤 2 – 添加新列

接下来，我们将通过基于 VendorID 和行程发生的月份对数据进行分组，归一化 total_amount 列。由于我们目前只有完整的时间戳，因此需要向 DataFrame 添加一列，从时间戳中提取月份：

from pyspark.sql.functions import *
df = df.withColumn('tpep_pickup_datetime', month(col('tpep_pickup_datetime')))

现在，我们可以执行归一化操作。准备好——这将是一项复杂的任务！我们将创建并使用一个用户自定义函数 (UDF)。虽然你可以使用 Spark 内置的本地函数来进行计算，但将归一化操作作为 UDF 是为了模拟你可能需要的任何复杂逻辑。与本章的主题相呼应——将 Spark 集成到机器学习工作流中，模型训练和预处理的重度计算通常通过自定义 UDF 来实现。这使得我们可以利用现有的 Python 代码和工具，而不牺牲性能。使用 Python 来表达特定的计算和逻辑通常比使用 Spark 内置函数要简单得多，包括以下操作：

加权平均
加权相关
指数移动平均

现在，进入步骤 3。

步骤 3 – 创建 UDF 以归一化列

在这个示例中，我们将归一化 total_amount 列。归一化的标准公式如下：

(值 - 平均值)÷标准差\text{{(值 - 平均值)}} \div \text{{标准差}}(值 - 平均值)÷标准差

注意
在继续之前，我要说明一下，这个示例是从 Julien Le Dem 和 Li Jin 进行的精彩网络研讨会中改编的，他们在该研讨会中介绍了这一功能。他们提供了比我更多的细节，强烈建议你有机会时观看这个研讨会。可以在这里找到：www.dremio.com/webinars/im…。

有趣的是，当你使用 PySpark 创建 UDF 时，它会在 Python 中执行你的函数，而不是在 Spark 运行的原生 Java/Scala 中。因此，UDF 的运行通常比内置的等效函数要慢，但编写起来更简单快捷。

UDF 有两种类型：

行操作函数：这些函数一次处理一行数据：
- lambda x: x + 1
- lambda date1, date2: (date1 – date2).years
分组操作函数：这些函数需要多行数据来计算：
- 例如按月份分组计算加权平均

在本例中，为展示利用 Arrow 的好处，我们将专注于分组 UDF 的情况。

在构建按月归一化数据的 UDF 时，我们需要一些样板代码来将多行数据打包成嵌套行，并在其中进行计算。由于这一点，性能会受到影响，因为 Spark 必须物化分组并将其转换为 Python 数据结构，以运行 UDF，这非常耗费资源。我们来试试这个示例；在查看 notebook 中的解决方案之前，尝试自己实现：

为了节省输入时间，先导入 PySpark SQL 类型函数：

from pyspark.sql.types import *

接下来，我们创建一个结构列来表示我们的嵌套行：

group_cols = ['VendorID', 'pickup_month']
non_group_cols = [col for col in df.columns if col not in group_cols]
s = StructType([f for f in df.schema.fields if f.name in non_group_cols])
cols = list([col(name) for name in non_group_cols])
df_norm = df.withColumn('values', struct(*cols))

现在，我们可以使用分组定义来聚合 DataFrame 中的值：

df_norm = (df_norm
           .groupBy(*group_cols)
           .agg(collect_list(df_norm.values)
           .alias('values')))

你可以看到有很多样板代码。接下来，我们需要定义 UDF 来执行归一化：

s2 = StructType(s.fields + [StructField('v', DoubleType())])
@udf(ArrayType(s2))
def normalize(values):
    v1 = pd.Series([r.total_amount for r in values])
    v1_norm = (v1 - v1.mean()) / v1.std()
    return [values[i] + (float(v1_norm[i]),) for i in range(0, len(values))]

我们快完成了！现在有了归一化函数，我们只需应用它，展开列，并删除我们仅用于计算的多余列：

df_norm = (df_norm.withColumn('new_values', normalize(df_norm.values))
           .drop('values')
           .withColumn('new_values', explode(col('new_values'))))

for c in [f.name for f in s2.fields]:
    df_norm = df_norm.withColumn(c, col('new_values.{0}'.format(c)))

df_norm = df_norm.drop('new_values')
df_norm.show()

关键行是启动 Spark 执行计算并开始处理所有任务的地方。在此之前，它只是创建了一个执行计划，直到我们需要它展示结果为止。

准备好所有内容后，我们在第一行加上 %%time，然后运行它，看看需要多长时间：

CPU times: user 8.56 ms, sys: 11.9 ms, total: 20.4 ms
Wall time: 30.5 s

还不错吧？大约 30 秒。再次注意，CPU 计时显示 Python 进程只花了大约 20 毫秒来构建计划并将其发送给 JVM。30 秒的运行时间全部发生在 JVM 进程以及 Spark 启动的 Python 进程中，用于运行我们创建的 UDF。对于 350 万行的数据来说，这个简单计算花费 30 秒似乎有点长。那么，为什么耗时这么久呢？我们来看看：

我们必须打包和解包嵌套行以获取所需数据。
为了传递数据，进行了大量序列化/反序列化。默认情况下，Spark 使用 Python 的 pickle 协议。
计算仍然是标量的，因此我们为 Python 解释器和按行处理模式支付了开销，而不是按列处理模式。

我们能否在清理样板代码的同时提高性能？当然可以，不然我就不会用这个例子了！

步骤 4 – 抛弃步骤 3，改用矢量化 UDF

首先，我将向你展示矢量化计算的示例，然后我会深入解释其中的差异。那么，准备好了吗？这是矢量化 UDF 及其使用方法：

schema = StructType(df.schema.fields + [StructField('v', DoubleType())])

def vector_normalize(values):
    v1 = values.total_amount
    values['v'] = (v1 - v1.mean()) / v1.std()
    return values

group_columns = ['VendorID', 'pickup_month']
df_pandas_norm = df.groupby(*group_columns) \
    .applyInPandas(vector_normalize, schema=schema)

df_pandas_norm.show()

就是这么简单。这就是全部内容。你应该会看到与步骤 3 相同的输出，只是速度要快得多。到底快多少呢？加上我们神奇的 %%time 关键字，让我们一起来看看：

CPU times: user 6.65 ms, sys: 0 ns, total: 6.65 ms
Wall time: 2.48 s

现在是 2.48 秒，而不是 30.5 秒。计算时间减少了 91.86% ！那么，我们是怎么做到的？为什么运行速度会这么快？答案在于细节。

步骤 5 – 理解背后发生了什么

要理解为什么矢量化计算能带来如此巨大的好处，首先你需要了解 PySpark UDF 是如何计算的。图 9.7 展示了这个执行过程的简化图：

在执行 UDF 的过程中，运行在 Java/Scala 中的 Spark 执行器会将成批的行流式传输给 Python 工作进程。该进程只会使用一个 for 循环，逐行调用 UDF 并将结果作为另一批行返回。由于 Spark 执行器运行在与 Python 工作进程完全不同的运行环境中，你可以猜到其中一个主要的开销就是数据的序列化和反序列化，用来在不同的环境之间来回传输数据。除此之外，由于这仍然是标量计算，逐行循环并不是执行计算的最高效方式。

然而，如果我们能够利用 pandas 中实现的矢量化列式计算，并消除序列化/反序列化过程，如图 9.8 所示，最终我们可以获得更好的性能：

在 Spark 中还有其他方式可以利用 pandas 的矢量化计算：

使用 mapInPandas：可以用 mapInPandas 代替 applyInPandas，通过迭代器而不是基于分组的操作来应用函数，比如根据某些条件过滤行。
简单的 pandas UDF：当不按某个值对数据进行分组时，可以使用简单的 pandas UDF 来执行矢量化的逐元素操作，比如将一列与另一列相乘。这是一种比传统 Spark UDF 更快的替代方案。
分组聚合：对于按组聚合的操作，例如计算某列的平均值或标准差，可以通过将 pandas_udf 与窗口函数结合使用来加速，而不是使用传统的 udf 装饰器。

更多信息可以在 Spark 的 pandas 和 Arrow Python 文档中找到：Spark Python 文档。希望前面的示例可以帮助你了解如何利用矢量化列式计算来优化 UDF 以及任何你希望使用 Spark 进行扩展的机器学习工作流。

步骤 6 – 从我们的努力中获益

最终，对于数据科学家来说，Jupyter 和 Apache Spark 是一个常见的开发环境，而它又是另一个可以利用 Arrow 来提升性能的工具。正如我们所看到的，从加载数据文件到执行计算的每一步，都有可能从 Arrow 提供的通用内存格式中受益。现在，让我们回到数据处理流程的前一步。在使用 Spark 对数据进行预处理之前，你首先需要这些数据集！为此，我们可以从 Hugging Face 获得一些帮助。

Facehuggers 植入数据

Hugging Face（huggingface.co）是一家总部位于纽约的公司，最初开发了一个针对青少年的聊天机器人，后来开源了该聊天机器人的机器学习模型。在开源模型之后，Hugging Face 转变了他们的商业模式，专注于构建一个完整的机器学习平台。作为一个平台，他们提供了以下服务：

基于 Git 的机器学习模型仓库
包含文本、图像、视频和音频的数据集
提供小规模机器学习演示的 Web 应用程序

此外，他们还提供了多种用于机器学习工作流的库，如 transformers、tokenizers、accelerators 以及（重要的）dataset 库。我们现在要重点关注的就是这个 dataset 库。

在 Hugging Face Hub 上，社区的开发者和数据科学家管理了大量数据集。其中许多数据集是公开的，而另一些则是私有的，访问这些私有数据集需要订阅 Hugging Face 的付费服务。你可以在他们的浏览器中心找到这个数据集库，按主题、许可、语言等条件搜索数据集：huggingface.co/datasets。现在，我想特别强调他们的 Python 库，名为 datasets。这个库提供了多种数据处理功能，帮助你准备好数据用于机器学习训练、验证等。但是与我们目前讨论最相关的是，datasets 库大量使用了 Apache Arrow！让我们安装这个库，下载一个数据集，然后看看你能否发现 Arrow 的使用之处。

注意！
你可以在 Hugging Face 网站上找到 datasets 库的完整文档：huggingface.co/docs/datase…。我们只会覆盖文档中的一小部分内容，来展示与 Arrow 的集成。

设置环境

我强烈建议使用 venv 虚拟环境来避免依赖问题，并且便于后续清理。不过，这并不是强制的：

$ python -m venv .venv  # 创建虚拟环境
$ source .venv/bin/activate  # 激活环境

如果稍后你想重用终端，但不再使用虚拟环境，只需使用 deactivate 函数：

(.venv) $ deactivate  # 停用虚拟环境

无论你是否使用虚拟环境，安装库都非常简单：

$ pip install datasets

你也可以使用 conda 来安装：

$ conda install -c huggingface -c conda-forge datasets

安装好库之后，我们可以轻松从 Hugging Face Hub 加载数据集。继续跟随我——很快一切都会变得清晰！让我们加载一个简单的数据集。首先，请确保你有至少 30 GB 的可用空间来下载它，然后运行以下 Python 代码：

>>> from datasets import load_dataset
>>> wiki = load_dataset('wikipedia', '20220301.en', split='train', trust_remote_code=True)

加载数据集分片：
首次运行时，可能需要一些时间，因为需要下载 30 GB 的 Parquet 文件！请耐心等待一段时间。完成后，你将拥有一部分 Wikipedia 的本地副本，接着我们可以检查一些关于数据集的信息：

>>> wiki.num_columns
4
>>> wiki.num_rows
6458670
>>> print(f"size: {wiki.dataset_size >> 30} GB")
size: 18 GB
>>> type(wiki)
<class 'datasets.arrow_dataset.Dataset'>

注意到了吗？我们的数据集的类名是 arrow_dataset.Dataset！你现在能猜出 Hugging Face 的 datasets 库是如何使用 Arrow 的吗？我们可以通过检查 info 属性进一步深入了解：

>>> wiki.info
DatasetInfo(description='Wikipedia dataset containing....

如果你查看完整输出，你会看到它提到了数据集的一些 Parquet 文件。但我们还需要看另一个属性来弄清一切：

>>> wiki.cache_files
[{'filename': '/home/matt/.cache/huggingface/datasets/wikipedia/20220301.en/2.0.0/....../wikipedia-train-00000-of-00041.arrow'}, {'filename':......

所有缓存文件都有 .arrow 扩展名！为什么会这样？回想一下第3章《格式和内存处理》，当时我们讨论了将 Arrow IPC 格式写入文件的好处。你应该记得我们介绍了一个重要概念——内存映射，它允许我们高效处理大于内存的数据文件，只有在我们试图访问某列的某个部分时才分配内存。这正是 Hugging Face 数据集缓存系统所利用的。我们可以通过检查加载数据集时的实际内存使用情况来确认这一点。

提示在代码示例的输出中，你可以看到默认的缓存目录是 ~/.cache/huggingface/datasets。当然，这可能不是你想要缓存数据的位置。如果想更改 Hugging Face 用于缓存 .arrow 文件的目录，有两种主要方法：设置 HF_DATASETS_CACHE 环境变量为一个目录路径，或者在 load_dataset 调用中添加 cache_dir="/path/to/cache" 参数。

通过检查资源使用情况证明优势

首先，我们安装 psutil 模块（我们在第 3 章中使用过它来演示内存使用情况）：

$ pip install psutil

现在，我们可以展示一些内存使用情况的数据：

>>> import os
>>> import psutil
>>> from datasets import load_dataset
>>> mem_before = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
>>> wiki = load_dataset('wikipedia', '20220301.en', split='train')
Loading dataset shards.....
>>> mem_after = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
>>> print(f"Mem Used: {(mem_after - mem_before)} MB")
Mem Used: 428.87890625 MB

加载 18 GB 的数据集仅占用我们 Python 进程中 400 多 MB 的实际内存。这是因为数据集通过缓存从磁盘进行内存映射，而不是一次性将整个数据集加载到内存中。由于这个机制，我们可以非常快速地迭代数据集的批次，我可以通过 timeit 模块证明这一点：

>>> import timeit
>>> s = """batch_size = 1000
... for batch in wiki.iter(batch_size):
...    pass
... """
>>> elapsed_time = timeit.timeit(stmt=s, number=1, globals=globals())
>>> print(f"Time to iterate {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 18 GB dataset: 45.2 sec, 3.3 Gb/s

第一行代码使用 timeit 模块来测量执行我们创建的字符串中的 Python 代码所需的时间。在这个示例中，我们只运行一次来获取一个测量值，但你可以增加 number 参数的值，让它连续运行多次，然后将秒数除以运行次数以获得平均运行时间。第二行是我们的输出：我们能够在约 45 秒内完全迭代一个 18 GB 的数据集，速度为 3.3 GB/秒！

虽然 Hugging Face 的 datasets 库提供了大量工具来处理、操作和流式传输数据，但如果你愿意，也可以通过相应的 to_pandas 和 to_polars 方法将 Dataset 对象转换为更熟悉的 pandas 和 Polars DataFrame。

一旦我们的 Dataset 对象利用 Arrow 表格的零拷贝优势进行内存映射，我们只需要将这些数据提供给一些标准工具，然后就可以开始构建机器学习工作流了！

在标准的机器学习工具中使用 Arrow

随着 Python 机器学习工具和实用程序生态系统的迅速发展，一些框架已经成为构建训练和推理管道的事实标准。其中最流行的是 PyTorch 和 TensorFlow，它们与 Hugging Face 的集成相当紧密，还有许多基于它们构建的系统。TensorFlow 和 PyTorch 都是开源库，前者以 Apache License 2.0 许可发布，后者则使用 BSD-3 许可发布。

这两个框架的主要数据结构是张量（tensor）或 n 维数组。机器学习模型通常由多层计算组成，每一层都有一个张量作为输入，并输出一个张量到下一层。简而言之，你可以将张量描述如下（如图 9.9 所示）：

一维张量通常称为向量，例如 [1, 2, 3, 4]。
具有等长行的二维张量称为矩阵，并由形状定义。例如，形状为 (3, 2) 的矩阵可以表示为 [[1, 2, 3], [4, 5, 6]]。
你还可以有一个三维张量，通常称为立方体。

张量是这些概念的通用术语，并且不对维度数量进行上限限制。

鉴于本书并非专门讲解机器学习，我不会复述众多关于使用这些库的教程。重要的是，这些延迟加载的数据集可以被传递给这些框架，并利用 Apache Arrow 底层的优势。

进一步阅读
如果你确实想深入了解如何利用这些库来构建完整的管道，可以查看以下教程：

这些工具的共同点是都使用 Arrow 作为数据源，进行内存管理和互操作性。甚至 TensorFlow 本身的 I/O 库中也提供了对 Arrow 数据集的底层支持：tfio.arrow.ArrowDataset 和 tfio.arrow.ArrowStreamDataset。结合许多库对 pandas 和 Polars 的广泛支持（包括与 Arrow 的零拷贝转换），Arrow 仍然是进出机器学习工作流和数据处理工具的最佳方式之一。当然，当你想到机器学习和人工智能时，首先想到的现在通常是利用 GPU 来提升速度。

总结

总而言之，任何 AI 或机器学习管道的核心都是数据处理。这意味着，任何能够简化数据访问、处理并在工具之间传递的工具对构建模块化和可组合的系统都至关重要。通过利用内存映射、DLPack 协议、Arrow C 数据/设备接口以及张量的规范扩展类型，越来越多的库和框架开始采用并受益于 Arrow 格式及其实现。

到此为止，我们暂时不会继续深入开发 Arrow，而是转向探讨 Arrow 在实际应用中的使用示例。下一章名为《Powered by Apache Arrow》，我们将探讨现有的由 Arrow 驱动的应用程序和用例，看看它们是如何使用 Arrow 的。我想特别介绍一些创新项目，它们在创意方面使用了 Arrow。也许你已经熟悉其中一个或多个用例，或者我将为你介绍你新的最爱工具。

只有一种方法可以发现答案——继续阅读吧！

使用 Apache Arrow 进行内存分析——在机器学习工作流中使用 Arrow