企业应用的需求日益增长。智能体 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：

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

• UDF 因回退而产生往返数据传输开销。

• UDF 放弃了潜在的计算加速，如果它是用 GPU 加速原语编写的话。

加速 UDF

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

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

| 方法 | 复杂度 | 适用场景 |
| --- | --- | --- |
| SQL | 低，最高抽象层级 | 简单的关系逻辑，例如小型或内联函数 |
| cuDF | 中等，对执行有更多灵活性和控制力 | 大多数情况，例如字符串解析、嵌套类型等中等复杂度的逻辑 |
| CUDA（NVIDIA 并行计算平台和编程模型） | 高，最大灵活性和内核级控制 | 最复杂或性能关键的逻辑 |

为了使这些方法更加具体，考虑一个简单的 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 表达式替换 UDF：

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 函数。该函数随后组合 Java API（应用程序编程接口） for cuDF operations 中的操作来并行处理列。这种方法为更复杂的处理提供了更丰富的 API。

使用 CUDA 重写 UDF

对于无法从现有 cuDF 操作高效组合的高度复杂逻辑，RapidsUDF 可以通过 Java Native Interface (JNI) 调用自定义 CUDA 内核。为简洁起见，我们省略了代码，但 CUDA 转换的示例可以在 cuDF examples for Spark UDFs 中找到。

使用 AI 智能体重写以提高效率

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

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

可用技能

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

图 2. 四项技能引导智能体完成测试、转换、基准测试和优化，每个阶段都有验证驱动的反馈循环

脚手架与设计

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

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

性能结果

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

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

开始使用

将 UDF 转换为 GPU 实现可以带来显著的加速，并有助于在 GPU 加速管道中整体降低成本。借助智能体技能，重写过程可以变得既简单又可靠。

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

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

如需深入了解该工作流程以及构建过程中的实践经验，请观看 Data + AI Summit 会议：Coding with AI Agents to Supercharge Your Spark UDFs with GPU Acceleration。

如需了解更多关于 Apache Spark 的 GPU 加速，请参阅适用于 Apache Spark 的 cuDF 插件入门指南，以在您的 Spark 管道中解锁更优的性能和 TCO。访问我们的网页以了解更多关于 NVIDIA cuDF 的信息：GPU-Accelerated Structured Data Processing，或者您可以阅读文档以获取有关 NVIDIA cuDF plugin for Apache Spark 的更多信息。

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：NVIDIA Technical Blog

原创地址：https://developer.nvidia.com/blog/rewrite-apache-spark-udfs-gpu-acceleration-ai-agents/