如何利用 AI Agent 重写 Apache Spark UDF 以实现 GPU 加速 | NVIDIA 技术博客

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

满足这一需求需要更高的数据吞吐量，且成本不能成比例增加。NVIDIA cuDF 通过利用 GPU 的大规模并行性和高带宽内存进行数据分析来应对这一挑战。特别是，适用于 Apache Spark 的 NVIDIA cuDF 插件将这种加速带到了 Apache Spark（行业领先的大数据处理和分析平台），从而在主要工作负载中实现更好的 TCO 性能。

用户自定义函数 (UDF) 长期以来一直是加速查询引擎的难题，由于其不透明的性质，会引入运行时开销。虽然 UDF 通常可以用更高性能的代码替换，但重写 UDF 长期以来一直是一个费时费力的过程，需要仔细的验证。

本文介绍了如何使用智能体在几分钟内将 UDF 重写为 GPU 加速的等效代码，从而在 Apache Spark 作业中实现 22x+ 的加速。我们将介绍适用于 Apache Spark 的 cuDF 技能，这是一套旨在快速突破 UDF 瓶颈的技能，涵盖从测试到基准测试再到集成的端到端开发。

UDF 的挑战

cuDF 能够实现无缝的 GPU 加速，部分归功于 Apache Spark 的声明式编程接口。SQL / Dataframe 抽象非常强大，因为它们将查询应该产生的结果与系统应该如何执行它分离开来。这使得 cuDF 插件能够利用 GPU 的强大 SIMD 能力，在不更改用户代码的情况下，透明地并行化行级操作。

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

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

• UDF 因回退而产生往返数据传输开销。
• 如果使用 GPU 加速原语编写，UDF 就能获得潜在的计算加速，而现在却失去了这一机会。

加速 UDF

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

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

为了使这些方法具体化，考虑一个将温度从摄氏度转换为华氏度的简单 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：