AWS Glue 6.0 实测：原生 Arrow UDF 优化 PySpark 性能

AWS Prescriptive Guidance 主题 优化用户定义函数 解释说，PySpark 中的 Python UDF 会降低性能，因为数据必须在 Java 虚拟机（JVM）和 Python 运行时之间进行序列化。它建议按优先顺序：使用内置 Spark SQL 函数，用 Scala 或 Java 实现 UDF，或对向量化工作负载使用基于 Apache Arrow 的 UDF（也称为向量化 UDF 或 pandas UDF）。

该指导基于一个简洁的示例：操作 a + b 作用于 1000 万行，比较标准 Python UDF、向量化 UDF 和 Spark SQL。在本文中，我们扩展该示例来回答我们从客户那里听到的几个实际问题：

• 随着数据量从 1000 万行增长到 1 亿行，该排序是否仍然成立？
• 它是否适用于计算更密集的操作，而不仅仅是 a + b？
• AWS Glue 6.0（Apache Spark 4.1.1）增加了两个较新的基于 Arrow 的选项：Arrow 优化 Python UDF（useArrow=True）和原生 Arrow UDF（arrow_udf，来自 Apache Spark 4.1.0）。它们相对于向量化（pandas）UDF 处于什么位置？
• 我们能否将差异归因于数据传输、计算，还是两者兼有？

我们将每个变体作为 AWS Glue 6.0 作业运行，并分享结果、方法以及一个受控实验，该实验阐明了向量化 UDF 表现如此的原因。

解决方案概述

我们将一个参数化基准测试框架作为 AWS Glue 6.0 作业运行。该作业生成一次 Parquet 数据集，将其写入 Amazon Simple Storage Service（Amazon S3），然后针对相同的文件运行每个 UDF 变体，以便每种方法读取完全相同的数据。我们在三种数据量（1000 万、5000 万和 1 亿行，其中 M = 百万）和三种每行成本递增的工作负载上进行了测量。

下图展示了基准测试的工作流程。

图 1：将每个 UDF 变体作为 AWS Glue 6.0 作业运行的基准测试工作流程

高层步骤如下：

1. 设置必备的 AWS Identity and Access Management（IAM）角色和 Amazon Simple Storage Service（Amazon S3）存储桶。
2. 部署并运行 AWS Glue 6.0 作业，该作业生成数据集并运行各个变体。
3. 查看结果，结果输出到作业日志，并可选择写回 Amazon S3。

我们比较的方法

所有变体计算相同的结果。它们的区别在于数据如何跨越 JVM–Python 边界以及每行逻辑如何运行。

向量化 UDF 和 Arrow 优化 Python UDF 都使用 Apache Arrow 进行传输，因此两者在物理计划中都显示 ArrowEvalPython。区别在于计算：

• 一个 向量化 UDF（pandas UDF）接收一个 pandas.Series（一批行），因此计算可以用 NumPy 表达，并在原生代码中对整个批次运行。它在 Arrow 之上涉及 pandas 转换。
• 一个 Arrow 优化 Python UDF是标量 UDF：函数每次调用接收一个值。Arrow 使传输高效，但 Python 代码仍然一次运行一行。
• 一个 原生 Arrow UDF（arrow_udf 装饰器，在 Apache Spark 4.1 中引入并在 AWS Glue 6.0 中可用）接收并返回一个 pyarrow.Array，并使用 PyArrow 计算内核进行计算。它直接操作 Arrow 数据，无需向量化 UDF 所执行的 pandas 转换。

下图总结了每个变体如何跨越边界。

图 2：每个 UDF 变体如何跨越 JVM 到 Python 边界

代码差异反映了这一区别。注册形式相似，但计算不同：

# Arrow-optimized Python UDF: scalar function, Arrow transport
op = udf(lambda a, b: math.sqrt(a*a + b*b) * math.sin(a), DoubleType(), useArrow=True)

# vectorized UDF (pandas UDF): batch function, Arrow transport + NumPy computation
@pandas_udf(DoubleType())
def op(a, b):                       # a, b are pandas.Series (a batch of rows)
    return np.sqrt(a*a + b*b) * np.sin(a)

# native Arrow UDF: batch function, Arrow transport + PyArrow compute (no pandas)
@arrow_udf("double")
def op(a, b):                       # a, b are pyarrow.Array (a batch of rows)
    return pc.multiply(pc.sqrt(pc.add(pc.multiply(a, a), pc.multiply(b, b))), pc.sin(a))

环境与设置

• 运行时： AWS Glue 6.0（Apache Spark 4.1.1、Python 3.13、Scala 2.13），10 × G.1X 执行器加上驱动程序，位于美国东部（弗吉尼亚北部）AWS 区域（us-east-1）。
• 数据： 一个包含两个 bigint 列（a、b）的 Parquet 数据集，写入 Amazon S3 一次，并由每个变体读回，以便所有方法读取完全相同的文件。
• 数据量：1000 万、5000 万和 1 亿行。 每次运行结合两个维度：工作负载（每行操作）和 查询形态（该操作如何聚合）。三种工作负载中的每一种都在两种查询形态下运行，针对每种方法和数据量。

工作负载（每行操作）

我们选择了三种成本递增的操作，以展示当每行工作从“几乎为零”转向“仅 Python”时结果如何变化。

这种递进很重要。arith 隔离了传输成本。mathheavy 增加了真实计算，因此向量化可能有帮助。全部五种方法都适用，包括原生 Arrow UDF，因为数学运算可以映射到 PyArrow 计算内核上。pyheavy 是相反的极端：一个既不能映射到 SQL 也不能映射到 PyArrow 计算内核的迭代循环，因此 Spark SQL 和原生 Arrow UDF 不适用，只剩下基于 Python 的 UDF。

查询形态（操作如何聚合）

每个工作负载都在两种查询形态下测量，因为它们在引擎中走不同的执行路径：

• 流式： sum(op(a, b)) — 对每一行进行全扫描，归约为单个值。这以最小的聚合开销隔离出每行的成本。
• 有状态： GROUP BY cast(op(a, b) AS bigint) % 1,000,000 — 先进行全扫描，然后进行高基数聚合（约一百万个分组），这会增加一次 shuffle 和一个大型哈希表。

两种形态都适用于全部三个工作负载。为保持表格可读，以下结果报告的是流式形态。我们也测量了有状态形态，发现排名相同，代表性有状态数据在结果部分给出。

方法

有两个决策对结果产生了实质性影响。

预热与 N 次取最佳。每个变体的首次运行主要受集群、JVM 和 Python worker 启动的影响。我们运行一次预热迭代并丢弃，然后运行 N 次计时迭代，并报告最佳值（以及平均值）。

每次迭代都重新读取。我们最初的尝试复用了缓存的 DataFrame，而 Spark 缓存了 Parquet 块并为标量count()裁剪了计划。结果 UDF 的计时没有随数据规模扩展，Spark SQL 看起来比 UDF 还慢。这与预期结果相反。我们通过以下方式纠正：在每次计时迭代时清除缓存并重新读取 Parquet，以及使用sum(...)，使优化器无法裁剪每行的工作。

def bench(name, fn, repeat, warmup, pre=None):
    for _ in range(warmup):
        if pre: pre()
        fn()                     # discarded
    times = []
    for _ in range(repeat):
        if pre: pre()            # e.g. spark.catalog.clearCache()
        t = time.perf_counter(); fn(); times.append(time.perf_counter() - t)
    return {"best_sec": min(times), "mean_sec": sum(times) / len(times)}

经过这一纠正后，排名与 AWS Prescriptive Guidance 的建议一致，并在各数据量下保持一致。简而言之：如果一个 UDF 基准测试不随行数扩展，那它很可能测量的是缓存而不是计算。

结果

以下表格报告每个工作负载的流式查询形态。我们也验证了有状态形态，它产生相同的排名。例如，在mathheavy的 100M 行下，有状态顺序与流式一致（Spark SQL 1.59 秒，向量化 UDF 4.22 秒，Arrow 优化 Python UDF 11.02 秒，标准 Python UDF 14.33 秒）。后面的pyheavy表格并排展示了流式和有状态列。有状态数据比流式略慢，因为增加了 shuffle 和哈希聚合，但各方法的相对顺序不变。

工作负载 1：arith (a + b)

七次迭代取最佳，不含预热，流式：

推荐排名成立，且向量化 UDF 相对于标准 Python UDF 的优势随数据量扩大（从 10M 时约 2.2× 到 100M 时 3.7×），因为标准 Python UDF 的每行序列化成本随行数增长。对于这个简单操作，原生 Arrow UDF 和向量化 UDF 接近，因为没有太多可向量化的计算。它们相对标量方法的优势主要来自列式传输。

工作负载 2：mathheavy

7 次迭代取最佳，流式：

在 UDF 方法中，原生 Arrow UDF 最快。它在 100M 时比向量化 UDF 快约 18%（3.13 秒对 3.83 秒），因为它直接在 Arrow 数据上计算，避免了 pandas 转换。计时在 10M、50M 和 100M 之间大致线性扩展，排名在每个数据量下都稳定。下图以双对数坐标展示了各方法随数据量的变化。

图 3：mathheavy 工作负载上各方法在不同数据量下的运行时间，双对数坐标

工作负载 3：pyheavy（无 SQL 等价物）

此工作负载运行逐行迭代循环，没有闭式等价物。对于每一行，它从a和b推导出一个整数，并重复应用 Collatz 步骤（偶数则减半，否则3n + 1），直到该值达到 1，统计步数。为使每行运行时间可预测，循环是有界的：即使尚未达到 1，也最多在 1,000 次迭代后停止，因此没有哪一行能主导基准测试。

def collatz_steps(a, b):             # a, b are scalars (one row per call)
    n = int((a ^ b) % 100000) + 1
    steps = 0
    while n != 1 and steps < 1000:   # bounded: at most 1,000 iterations per row
        n = n // 2 if n % 2 == 0 else 3 * n + 1
        steps += 1
    return steps

这种数据相关的控制流无法表示为 SQL 标量表达式，也无法用 PyArrow 计算内核表达，因此 Spark SQL 和原生 Arrow UDF 都不适用，只比较基于 Python 的 UDF 方法。该表展示了两种查询形态在 10M、50M 和 100M 行下的结果，以说明排名在不同形态和规模下一致：

我们报告pyheavy最多到 100M 行。由于 Spark SQL 和原生 Arrow UDF 在此不适用，仅剩的选择是基于 Python 的 UDF，它们的计时接近线性扩展，且到 100M 时其相对顺序在两种形态下都已稳定。

此工作负载代表 UDF 不可避免的情况。示例包括迭代算法，或调用没有 SQL 对应物、也无法映射到 PyArrow 计算内核的 Python 库。这里向量化 UDF 是最有效的可用选项：在 10M 时比标准 Python UDF 快约 2.2×，到 100M 时扩大到约 3.8×。

受控实验：传输还是计算？

向量化 UDF 始终优于 Arrow 优化的 Python UDF。为了测试这种差异是来自传输（两者都使用 Arrow）还是计算（只有向量化 UDF 在批次上使用 NumPy），我们使用 RDD.mapPartitions 构建了缺失的组合（NumPy 计算搭配 pickle 传输，不使用 Arrow）。

图 4：传输与计算作为 2×2 实验进行对比

NumPy 搭配 pickle 的变体是所有变体中最慢的（在 1 亿行时为 21.07 秒，比标准 Python UDF 的 13.82 秒更慢）。原因是强制使用 pickle 传输需要离开 DataFrame API 转而使用 RDD.mapPartitions，这会同时改变两件事。它会在构建 NumPy 数组之前将每一行物化为 Python 对象，并且绕过了向量化 UDF 所受益的优化 DataFrame 执行引擎。

实际结论是，向量化 UDF 的收益并非孤立地来自 Arrow 传输。而是 Apache Arrow 集成到向量化 DataFrame 执行路径中。该路径结合了列式传输和 NumPy 计算。在该路径之外仅复现 NumPy 部分并不能获得同样的收益。

代码对比：Arrow 优化的 Python UDF 与向量化 UDF

这两种方法经常被混淆：两者都使用 Apache Arrow 进行传输，并且都在物理计划中显示 ArrowEvalPython，因此值得将它们并排查看。以下两个定义计算相同的 mathheavy 操作。注册方式几乎相同。区别在于函数体及其参数类型。

以下代码是 Arrow 优化的 Python UDF（useArrow=True），其中函数接收标量值并且每行运行一次：

import math
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

def op(a, b):                        # a, b are scalars (one row per call)
    x = float(a % 100000) + 1.0
    y = float(b % 100000) + 1.0
    return math.sqrt(x * x + y * y) * math.sin(x) + math.log(y)

# Arrow transport, scalar Python computation
op_arrow = udf(op, DoubleType(), useArrow=True)
spark.udf.register("op_arrow", op_arrow)

df.selectExpr("sum(op_arrow(a, b)) AS r").collect()

以下代码是向量化 UDF（pandas UDF），其中函数接收一个 pandas.Series（一个批次）并使用 NumPy 在整个批次上进行计算：

import numpy as np
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType

@pandas_udf(DoubleType())
def op_pandas(a, b):                 # a, b are pandas.Series (a batch of rows)
    x = (a % 100000).astype("float64") + 1.0
    y = (b % 100000).astype("float64") + 1.0
    return np.sqrt(x * x + y * y) * np.sin(x) + np.log(y)   # NumPy, whole batch

spark.udf.register("op_pandas", op_pandas)

df.selectExpr("sum(op_pandas(a, b)) AS r").collect()

以下代码是原生 Arrow UDF（arrow_udf），其中函数接收并返回一个 pyarrow.Array 并使用 PyArrow 计算内核进行计算，直接操作 Arrow 数据而无需 pandas 转换：

import pyarrow as pa
import pyarrow.compute as pc
from pyspark.sql.functions import arrow_udf

@arrow_udf("double")
def op_native(a, b):                 # a, b are pyarrow.Array (a batch of rows)
    x = pc.add(pc.cast(pc.subtract(a, pc.multiply(pc.divide(a, 100000), 100000)), pa.float64()), 1.0)
    y = pc.add(pc.cast(pc.subtract(b, pc.multiply(pc.divide(b, 100000), 100000)), pa.float64()), 1.0)
    return pc.add(pc.multiply(pc.sqrt(pc.add(pc.multiply(x, x), pc.multiply(y, y))), pc.sin(x)), pc.ln(y))

spark.udf.register("op_native", op_native)

df.selectExpr("sum(op_native(a, b)) AS r").collect()

Arrow 优化的 Python UDF 和向量化 UDF 在两行关键代码上有所不同：

• 参数类型 – 在 Arrow 优化的 UDF 中，a 和 b 是标量，因此函数体使用 math.sqrt 和 math.sin 并且每行运行一次。在向量化 UDF 中，a 和 b 是 pandas.Series，因此函数体使用 np.sqrt 和 np.sin 并且每批次运行一次。
• 计算运行的位置 – Arrow 优化的 UDF 仍然逐行执行 Python。向量化 UDF 在 NumPy 的原生代码中执行整个批次。

计算运行位置的这种差异是在规模化时区分这三者的关键。在 mathheavy 工作负载的 1 亿行上，Arrow 优化的 UDF 运行了 11.01 秒，向量化 UDF 为 3.83 秒，而原生 Arrow UDF 最快，为 3.13 秒。原生 Arrow UDF 略微领先，因为它通过 PyArrow 计算内核向量化计算，而没有向量化 UDF 所执行的 pandas 转换。但是，它要求逻辑能够用这些内核表达，而向量化 UDF 仍然是更灵活的选择。

建议

这些结果与 AWS Prescriptive Guidance 的建议一致，并补充了几点实用建议：

1. 当操作可以用 SQL 表达时，优先使用内置 Spark SQL 函数 – 计算在 JVM 上运行，并且优势随数据量增长而增大。
2. 当必须使用 UDF 且逻辑可以用 PyArrow 计算内核表达时，考虑使用原生 Arrow UDF（arrow_udf）– 在我们的测试中，它是最快的 UDF 方法，因为它在 Arrow 数据上向量化计算而无需 pandas 转换。它在 AWS Glue 6.0（Apache Spark 4.1）中可用。
3. 否则，使用向量化 UDF（pandas UDF）– 在我们的测试中，它比标准 Python UDF 快 2.2–3.8 倍。当逻辑用 pandas 或 NumPy 比用 PyArrow 计算内核更容易表达时，它是更灵活的选择。
4. Arrow 优化的 Python UDF（useArrow=True）改善了传输但没有改善计算 – 它们比标准 Python UDF 更快，但当每行工作量不可忽视时，向量化或原生 Arrow UDF 更快。当你希望对现有标量 UDF 做最小改动时，考虑使用 Arrow 优化的 Python UDF。
5. 留在 DataFrame API 中以保持 Arrow 的收益 – 正如受控实验所示，加速并非仅来自 Arrow 序列化。它来自集成到向量化 DataFrame 执行路径中的 Arrow。通过 RDD.mapPartitions 应用相同的 NumPy 逻辑是所有方法中最慢的，因为它离开了该路径。应在 DataFrame API 内将向量化逻辑实现为向量化 UDF 或原生 Arrow UDF，而不是降级到 RDD API。
6. 对于迭代式或依赖数据的逻辑（ pyheavy 情况），向量化 UDF 是备选方案 – 当操作既不能映射到 SQL 也不能映射到 PyArrow 计算内核时，Spark SQL 和原生 Arrow UDF 都不可用，而在剩余的 Python UDF 方法中，向量化 UDF 是最高效的。
7. 判断重写是否值得投入 – 将标量 UDF 迁移到向量化 UDF 或原生 Arrow UDF 是较小的代码改动，但仍需编写和测试。根据经验，当每行工作量不可忽视且数据量大或作业频繁运行时，这是值得的。加速（以及每秒成本节省）随行数和作业运行频率而扩展，而重写是一次性成本。对于小型、不频繁数据集上的简单操作，标准 Python UDF 可能完全足够。对于每天运行、涉及数亿行的计算密集型转换，重写通常在前几次运行中就能收回成本。
8. 具体数字取决于你的集群 – 此处所有结果均使用 AWS Glue 6.0 上的固定配置（10 × G.1X worker）。排名各种方法的差异是每个UDF如何跨越JVM–Python边界的属性，预计不会随集群规模而变化。然而，绝对计时以及与替代方案的任何交叉点都会随工作节点数量、工作节点类型和数据布局而变化。请将这些数据视为相对行为的指南，并针对你自己的配置进行验证。
9. 方法会影响结果——预热处理和缓存管理改变了我们的排名。一个UDF基准测试，如果其计时不随行数扩展，那很可能是在测量缓存。

结论

UDF策略的选择对运行时和成本都有可衡量且反复出现的影响。在我们100百万行的计算密集型工作负载上，最快的UDF方法（原生Arrow UDF，3.13秒）与标准Python UDF（13.82秒）之间的差距约为4.4倍，且输出完全相同。由于AWS Glue按秒计费，这一差异直接转化为每次作业运行的开支。

这些结果扩展了AWS Prescriptive Guidance的建议，涵盖从10百万行到100百万行的容量扩展、计算密集型工作负载，以及AWS Glue 6.0（Apache Spark 4.1）中可用的较新基于Arrow的UDF选项。作为方法说明，预热处理和缓存管理对排名有实质性影响。一个UDF基准测试，如果其计时不随行数扩展，那很可能是在测量缓存而非计算。

要了解更多，请参阅优化用户定义函数（AWS Prescriptive Guidance）以及AWS Glue文档。