如何使用 AI 代理将 Apache Spark UDF 重写为 GPU 加速版本

企业应用的要求日益提高。智能体式 AI 进一步加剧了对海量企业数据的需求，同时也带来了针对不断扩大的数据集进行频繁、复杂查询的需求。Enterprise applications have become increasingly demanding. Agentic AI has compounded the need for vast amounts of enterprise data with frequent, complex queries across increasingly large datasets.

满足这一需求，需要在不按比例增加成本的情况下提升数据吞吐量。NVIDIA cuDF 利用 GPU 强大的并行处理能力和高带宽内存进行数据分析，从而应对这一挑战。尤其是，Apache Spark 的 NVIDIA cuDF 插件将这一加速能力引入 Apache Spark——这款业界领先的大数据处理与分析平台——以提升主要工作负载的总体拥有成本（TCO）表现。Meeting that demand requires greater data throughput without a proportional increase in cost. NVIDIA cuDF addresses this challenge by leveraging the massive parallelism and high-bandwidth memory of GPUs for data analytics. Particularly, the NVIDIA cuDF plugin for Apache Spark brings this acceleration to Apache Spark, the industry-leading platform for big data processing and analytics, to achieve better TCO performance in major workloads.

用户定义函数（UDF）长期以来一直是加速查询引擎面临的难题，其不透明的特性会带来运行时开销。虽然通常可以用性能更好的代码替代 UDF，但重写 UDF 一直是项耗时费力的工作，还需要进行谨慎验证。User-defined functions (UDFs) have long been problematic for accelerated query engines, introducing runtime overheads due to their opaque nature. While UDFs can often be replaced with more performant code, rewriting UDFs has long been a laborious process requiring careful validation.

本文介绍如何利用智能体在几分钟内将 UDF 重写为 GPU 加速版本，使 Apache Spark 作业提速 22 倍以上。我们将介绍适用于 Apache Spark 的 cuDF 技能集，帮助快速突破 UDF 瓶颈，覆盖从测试、基准测试到集成的端到端开发流程。This post describes how to use agents to rewrite UDFs into GPU-accelerated equivalents within minutes, achieving 22x+ speedups in Apache Spark jobs. We’ll introduce the cuDF skills for Apache Spark, a set of skills to quickly break through the UDF bottleneck for end-to-end development, from testing, to benchmarking, through integration.

UDF 挑战The UDF challenge

cuDF 能够实现无缝 GPU 加速，部分得益于 Apache Spark 的声明式编程接口。SQL / DataFrame 抽象之所以强大，是因为它们将查询应生成什么结果与系统如何执行查询分离开来。这使 cuDF 插件能够利用 GPU 强大的 SIMD 能力，透明地并行化逐行操作，而无需更改用户代码。cuDF achieves seamless GPU acceleration in part due to Apache Spark’s declarative programming interface. The SQL / Dataframe abstractions are powerful because they separate what a query should produce from how the system should execute it. This enables the cuDF plugin to transparently parallelize row-wise operations using the massive SIMD capabilities of GPUs, without any change in user code.

然而，Apache Spark 也支持 UDF：用于扩展内置功能的自定义例程。UDF 虽然有助于实现特定业务逻辑，却会打破声明式抽象——它们包含任意用户代码，系统无法对其进行推理，也无法以通用方式将其并行化。在执行期间，GPU 加速器必须将数据从列式格式转换为按行格式，并退回 CPU 逐行执行 UDF：Apache Spark, however, also supports UDFs: custom routines to extend built-in functionality. While useful for implementing bespoke business logic, UDFs break the declarative abstraction—they contain arbitrary user code which the system cannot reason about or parallelize in any general sense. During execution, the GPU accelerator must convert the data from a columnar to row-wise format, and fall back to the CPU to execute the UDF row-by-row:

因此，在其他方面均已加速的查询中，UDF 有以下缺点：Thus in an otherwise accelerated query, UDFs have the following drawbacks:

• 由于回退机制，UDF 会产生往返数据传输开销。
  The UDF incurs a round-trip data transfer overhead due to fallback.
• 如果使用 GPU 加速的原语编写，UDF 本可以获得计算加速，但它放弃了这一潜在优势。
  The UDF forgoes potential compute speedup, had it been written using GPU-accelerated primitives.

加速 UDFsAccelerating UDFs

通过进行少量、有针对性的代码修改，并使用恰当的基本操作，许多 UDF 都可以改写为可由 GPU 加速的形式。改写 UDF 可以消除回退到加速器所带来的开销，性能提升最高可超过 20 倍。With a small, targeted code change, many UDFs can be rewritten into a GPU-accelerable form by using the right primitives. Rewriting UDFs eliminates the overhead associated with accelerator fallback, and can yield speedups of over 20x.

我们概述了重写 UDF 的三种方法：We outline three approaches to rewrite UDFs:

为了让这些方法更具体，我们来看一个简单的 Scala UDF，它将温度从摄氏度转换为华氏度：To make these approaches concrete, consider a simple Scala UDF that converts temperature from Celsius to Fahrenheit:

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 重写 UDFRewriting the UDF with SQL

对于简单的 UDF，最直接的解决方案是将 UDF 替换为等效的 Spark SQL 表达式：The most direct solution for simple UDFs is to replace the UDF with an equivalent Spark SQL expression:

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

这使该函数能够在 Spark 的表达式树中显示。在查询规划期间，cuDF 插件会将这些表达式替换为相应的 GPU 实现，使数据和执行过程完全保留在 GPU 上。This makes the function visible in Spark’s expression tree. During query planning, the cuDF plugin will replace these expressions with their corresponding GPU implementations, keeping the data and execution fully on the GPU.

使用 cuDF 重写 UDFRewriting the UDF with cuDF

当 SQL 原语不足或不适用时，UDF 可以改用 RapidsUDF 接口直接处理 GPU 列。该接口由 cuDF 插件提供，是一项约定，要求实现单个方法：evaluateColumnar：When SQL primitives are insufficient or undesirable, the UDF can instead operate directly on GPU columns using the RapidsUDF interface, a contract provided by the cuDF plugin. The contract requires a single method, 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，可用于执行更复杂的处理。The cuDF plugin recognizes instances of the RapidsUDF interface and passes a GPU-resident ColumnVector to the evaluateColumnar function. The function then composes operations from the Java API for cuDF operations to process the column in parallel. This approach provides a richer API for more complex processing.

使用 CUDA 重写 UDFRewriting the UDF with CUDA

对于无法通过现有 cuDF 操作高效组合实现的复杂逻辑，RapidsUDF 可以通过 Java Native Interface（JNI）调用自定义 CUDA 内核。为简洁起见，此处省略了代码；有关 CUDA 转换的示例，请参阅 cuDF 中针对 Spark UDF 的示例。For highly complex logic that cannot be efficiently composed from existing cuDF operations, a RapidsUDF can invoke a custom CUDA kernel through the Java Native Interface (JNI). We’ve elided the code for brevity, but examples of CUDA conversions can be found in the cuDF examples for Spark UDFs.

使用 AI 代理高效改写Efficient rewrites with AI agents

尽管 UDF 重写能够带来性能优势，但往往需要投入大量工程工作。它们要求仔细选择 API、验证语义并测量性能，而对于不太熟悉向量化编程的从业者来说，列式 API 可能并不熟悉。Despite the performance benefits, UDF rewrites can require significant engineering effort. They demand careful API selection, semantic validation, and performance measurement, while columnar APIs can be unfamiliar to practitioners with less experience in vectorized programming.

幸运的是，这项工作中的大部分都很适合由编码代理来完成。在设置适当的防护措施后，代理可以搜索不熟悉的 API、实现候选改写，并快速迭代，同时通过可执行测试和性能测量来判断每项更改是否可接受。Fortunately, much of this work is well suited to coding agents. With proper guardrails, an agent can search unfamiliar APIs, implement candidate rewrites, and iterate quickly while executable tests and performance measurements determine whether each change is acceptable.

可用技能Available skills

Apache Spark 的 cuDF 技能将重写工作拆分为四个聚焦的反馈循环。每项技能都定义了要执行的工作，以及进入下一阶段前所需的证据：The cuDF skills for Apache Spark decompose the rewrite into four focused feedback loops. Each skill defines both the work to perform and the evidence required before proceeding to the next stage:

图 2. 四项技能引导智能体完成测试、转换、基准测试和优化，并在每个阶段通过验证驱动的反馈循环进行迭代。Figure 2. Four skills guide agents through testing, conversion, benchmarking, and optimization, with verification-driven feedback loops at each stage

脚手架与设计Scaffolding and design

可靠的智能体执行始于可靠的开发环境。这些技能提供了完整的项目模板，其中包含端到端开发所需的 cuDF 和 Spark 专用脚手架，包括测试和基准测试框架、用于验证输出的编码例程以及构建配置。将这些知识编码到代码中后，智能体就可以专注于 UDF，而不必反复重建专用基础设施。Reliable agent execution begins with a reliable development environment. The skills provide a complete project template containing the cuDF- and Spark-specific scaffolding required for end-to-end development. This includes test and benchmark harnesses, encoded routines for output verification, and build configuration. By encoding this knowledge into code, the agent can focus on the UDF rather than repeatedly reconstructing specialized infrastructure.

这些技能还提供了使用 cuDF 进行 GPU 编程的最佳实践，包括内存管理和配置，可作为技巧与窍门的参考指南。为了充分发挥前沿模型的能力，智能体仍可根据上下文探索不同的实现方式，并依据测得的正确性和性能选择合适的方案。The skills also provide best practices for GPU programming with cuDF, including memory management and configuration. This serves as an advisory tips and tricks reference. To best leverage the power of frontier models, the agent remains free to explore different implementations based on surrounding context and select an approach based on measured correctness and performance.

性能结果Performance results

一家大型消费科技公司使用这一工作流，将生产数据管道中使用的 30 多个 UDF 转换为 GPU 加速实现。改写后的实现保留了原始 CPU UDF 的行为，同时消除了原本已由 GPU 加速的工作负载中成本高昂的回退操作。例如，以下基准测试展示了其中一个工作负载中的 UDF 在改写前后的性能表现：A major consumer technology company used this workflow to convert 30+ UDFs used in production data pipelines into GPU-accelerated implementations. The rewrites preserved the behavior of the original CPU UDFs while removing costly fallbacks from otherwise GPU-accelerated workloads. As an example, the following benchmark shows the resulting before-and-after performance for a UDF in one such workload:

这个特定的 UDF 影响了 3,000 多个作业，而在其中一个作业中，将其实现为支持 cuDF 的 RapidsUDF 后，速度提升了 22 倍。借助智能体，这些性能提升更容易实现，使从业者能够在数小时内反复迭代重写，而不必耗费数天或数周，同时不影响正确性。This particular UDF was affecting more than 3,000+ jobs, and in this job, saw a 22x speedup when implemented as a cuDF-enabled RapidsUDF. With the help of agents, these performance gains are much more attainable, enabling practitioners to iterate on rewrites on the order of hours rather than days or weeks, without compromising correctness.

开始使用Get started

将 UDF 转换为 GPU 实现，可以显著提升速度，并帮助降低 GPU 加速流水线的总体成本。借助智能体技能，重写过程可以变得既简单又可靠。Converting UDFs to GPU implementations can unlock major speedups and contribute to overall cost savings in GPU-accelerated pipelines. With agent skills, the rewrite process can be made both easy and reliable.

要开始使用，请在 GitHub 上安装适用于 Apache Spark 的 cuDF skills，可以直接安装，也可以使用 npx skill 安装程序：To get started, install the cuDF skills for Apache Spark on GitHub, either directly or using the npx skill installer:

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

如需深入了解这一工作流以及构建过程中的实践经验，请观看 Data + AI Summit 的专题演讲《借助 AI 智能体编写代码，通过 GPU 加速提升 Spark UDF 的性能》。For a deeper walkthrough of the workflow and practical lessons from building it, see the Data + AI Summit session, Coding with AI Agents to Supercharge Your Spark UDFs with GPU Acceleration.

要进一步了解 Apache Spark 的 GPU 加速，请按照 cuDF for Apache Spark 插件入门指南操作，提升 Spark 流水线的性能并降低总拥有成本（TCO）。访问我们的网页，进一步了解 NVIDIA cuDF：GPU 加速的结构化数据处理；您也可以阅读文档，了解有关 NVIDIA cuDF for Apache Spark 插件的更多信息。To learn more about GPU acceleration for Apache Spark, follow the cuDF plugin for Apache Spark getting-started guide to unlock better performance and TCO across your Spark pipelines. Visit our webpage to learn more about NVIDIA cuDF: GPU-Accelerated Structured Data Processing, or you can read the docs for more information on NVIDIA cuDF plugin for Apache Spark.