用 AI 代理将 Apache Spark UDF 重写为 GPU 加速实现

企业应用的需求日益增长。代理式 AI 进一步加剧了对海量企业数据的需求，需要在越来越大的数据集上频繁执行复杂查询。

要满足这一需求，就必须在不按比例增加成本的前提下提升数据吞吐量。NVIDIA cuDF 通过利用 GPU 的大规模并行能力和高带宽内存进行数据分析来应对这一挑战。特别是，NVIDIA cuDF plugin for Apache Spark 将这种加速带到了 Apache Spark——业界领先的大数据处理与分析平台——从而在主要工作负载中实现更优的 TCO 表现。

用户定义函数（UDF）长期以来一直是加速查询引擎的难题， 由于其不透明的特性而引入了运行时开销。虽然 UDF 通常可以用性能更高的代码来替代，但重写 UDF 长期以来一直是一项费力的工作，需要仔细验证。

本文介绍如何使用代理在几分钟内将 UDF 重写为 GPU 加速的等效实现，从而在 Apache Spark 作业中实现 22 倍以上的加速。我们将介绍 cuDF skills for Apache Spark，这是一组技能，可快速突破 UDF 瓶颈，覆盖从测试、基准测试到集成的端到端开发流程。

UDF 的挑战

cuDF 之所以能够实现无缝的 GPU 加速，部分原因在于 Apache Spark 的声明式编程接口。SQL / Dataframe 抽象之所以强大，是因为它们将查询应产生的结果与系统应如何执行查询分离开来。这使得 cuDF 插件能够利用 GPU 的大规模 SIMD 能力透明地并行化逐行操作，而无需更改任何用户代码。

然而，Apache Spark 也支持 UDF：用于扩展内置功能的自定义例程。虽然 UDF 对于实现定制业务逻辑很有用，但它们打破了声明式抽象——它们包含任意用户代码，系统无法对其进行推理，也无法以任何通用方式并行化。在执行过程中，GPU 加速器必须将数据从列式格式转换为逐行格式，并回退到 CPU 逐行执行 UDF：

图 1. CPU UDF 迫使原本驻留在 GPU 上的流水线将列式批次转换为行，在 CPU 上执行，然后再转换回列式格式

因此，在原本已加速的查询中，UDF 具有以下缺点：

• 由于回退，UDF 会产生往返数据传输开销。
• 如果 UDF 使用 GPU 加速原语编写，本可获得潜在的计算加速，但 UDF 放弃了这一点。

加速 UDF

通过少量有针对性的代码修改，许多 UDF 可以通过使用正确的原语重写为可 GPU 加速的形式。重写 UDF 消除了与加速器回退相关的开销，并且可以实现超过 20 倍的加速。

我们概述了重写 UDF 的三种方法：

表 1. 重写 Spark UDF 的 SQL、cuDF 和 CUDA 方法比较，按复杂度和最适合用例组织

为了使这些方法具体化，考虑一个简单的 Scala UDF，它将温度从摄氏度转换为华氏度：

final class ToFahrenheit
extends Function1[java.lang.Double, java.lang.Double]
with Serializable {
override def apply(celsius: java.lang.Double): java.lang.Double = {
if (celsius == null) null
else celsius * 9.0 / 5.0 + 32.0
}
}
spark.udf.register("to_fahrenheit", new ToFahrenheit)
spark.sql("""
SELECT
*,
to_fahrenheit(celsius) AS fahrenheit
FROM temperatures
""")

使用 SQL 重写 UDF

对于简单的 UDF，最直接的解决方案是将其替换为等效的 Spark SQL 表达式：

spark.sql("""
SELECT
*,
celsius * 9.0 / 5.0 + 32.0 AS fahrenheit
FROM temperatures
""")

这使得该函数在 Spark 的表达式树中可见。在查询规划期间，cuDF 插件会将这些表达式替换为对应的 GPU 实现，从而将数据和执行完全保留在 GPU 上。

使用 cuDF 重写 UDF

当 SQL 原语不足或不适用时，UDF 可以改为使用 RapidsUDF

接口直接操作 GPU 列，这是 cuDF 插件提供的一种契约。该契约要求实现单个方法 evaluateColumnar

：

import ai.rapids.cudf.{ColumnVector, Scalar}
import com.nvidia.spark.RapidsUDF
import scala.util.Using
final class ToFahrenheit
extends Function1[java.lang.Double, java.lang.Double]
with Serializable
with RapidsUDF {
override def apply(celsius: java.lang.Double): java.lang.Double = {
if (celsius == null) null
else celsius * 9.0 / 5.0 + 32.0
}
override def evaluateColumnar(
numRows: Int,
args: ColumnVector*): ColumnVector = {
Using.Manager { use =>
val scale = use(Scalar.fromDouble(9.0 / 5.0))
val offset = use(Scalar.fromDouble(32.0))
val scaled = use(args.head.mul(scale))
scaled.add(offset)
}.get
}
}

cuDF 插件会识别 RapidsUDF

接口的实例，并将驻留在 GPU 上的 ColumnVector

传递给 evaluateColumnar

函数。随后，该函数会组合来自 cuDF 的 Java API 操作的操作，以并行处理该列。这种方法为更复杂的处理提供了更丰富的 API。

使用 CUDA 重写 UDF

对于无法由现有 cuDF 操作高效组合而成的极其复杂的逻辑，RapidsUDF

可以通过 Java 本地接口（JNI） 调用自定义 CUDA 内核。为简洁起见，我们省略了代码，但 CUDA 转换的示例可在 Spark UDF 的 cuDF 示例 中找到。

使用 AI 代理进行高效重写

尽管性能有所提升，但 UDF 重写可能需要大量的工程投入。它们需要谨慎的 API 选择、语义验证和性能测量，而对于向量化编程经验较少的从业者来说，列式 API 可能并不熟悉。

幸运的是，这些工作中的大部分都非常适合由编码代理来完成。在适当的防护措施下，代理可以搜索不熟悉的 API、实现候选重写，并快速迭代，同时由可执行测试和性能测量来决定每次更改是否可接受。

可用技能

用于 Apache Spark 的 cuDF 技能将重写分解为四个聚焦的反馈循环。每项技能都定义了要执行的工作以及进入下一阶段之前所需的证据：

图 2。四项技能引导代理完成测试、转换、基准测试和优化，并在每个阶段设置由验证驱动的反馈循环

脚手架与设计

可靠的代理执行始于可靠的开发环境。这些技能提供了一个完整的项目模板，其中包含端到端开发所需的 cuDF 和 Spark 专用脚手架。这包括测试和基准测试框架、用于输出验证的编码例程以及构建配置。通过将这些知识编码到代码中，代理可以专注于 UDF，而不是反复重建专门的基础设施。

这些技能还提供了使用 cuDF 进行 GPU 编程的最佳实践，包括内存管理和配置。这可作为提示与技巧的参考。为了最好地利用前沿模型的能力，代理可以自由地根据周围上下文探索不同的实现，并根据测量到的正确性和性能选择一种方法。

性能结果

一家大型消费科技公司使用此工作流，将生产数据管道中使用的 30 多个 UDF 转换为 GPU 加速实现。重写保留了原始 CPU UDF 的行为，同时从原本已 GPU 加速的工作负载中移除了代价高昂的回退。例如，以下基准测试显示了其中一个此类工作负载中某个 UDF 的前后性能对比：

图 3。使用 cuDF 重写生产环境中基于 CPU 的 UDF，将算子时间从 274.55 小时降至 12.78 小时，在 6000 亿行的工作负载上实现了 22 倍加速

这个特定的 UDF 影响了 3000 多个作业，而在这个作业中，当实现为启用 cuDF 的 RapidsUDF 时，实现了 22 倍加速。借助智能体，这些性能提升变得更加容易实现，使从业者能够以小时而非数天或数周为周期迭代重写，同时不牺牲正确性。

开始使用

将 UDF 转换为 GPU 实现可以释放巨大的加速潜力，并有助于在 GPU 加速的流水线中整体节省成本。借助智能体技能，重写过程可以变得既简单又可靠。

要开始使用，请安装 GitHub 上适用于 Apache Spark 的 cuDF 技能，可以直接安装，或使用 npx

技能安装器：

npx skills add NVIDIA/cudf-spark --skill ‘*’

如需深入了解该工作流程的详细讲解以及构建过程中的实践经验，请观看 Data + AI Summit 会议：使用 AI 智能体编码，通过 GPU 加速为你的 Spark UDF 提速。

如需进一步了解 Apache Spark 的 GPU 加速，请参阅 Apache Spark 的 cuDF 插件入门指南，以在您的 Spark 管道中实现更优的性能和总体拥有成本（TCO）。访问我们的网页，了解更多关于 NVIDIA cuDF：GPU 加速的结构化数据处理 的信息，或者您可以阅读文档，获取关于 NVIDIA cuDF plugin for Apache Spark 的更多信息。