Rewriting Apache Spark UDFs for GPU Acceleration Using AI Agents and NVIDIA cuDF

Rewriting Apache Spark UDFs for GPU Acceleration Using AI Agents and NVIDIA cuDF NVIDIA cuDF plugin accelerates Apache Spark workloads, but CPU-bound user-defined functions (UDFs) break the declarative execution model and trigger slow fallbacks. By using specialized AI agent skills, developers can automate the conversion of UDFs to GPU-accelerated forms, achieving substantial performance gains. CPU UDFs force Spark to convert GPU columnar batches to row format and fall back to the CPU, causing round-trip transfers and eliminating GPU compute speedups. UDFs can be rewritten using three approaches depending on complexity: native Spark SQL expressions, cuDF columnar operations via the RapidsUDF interface, or custom CUDA kernels via JNI. The cuDF skills for Apache Spark provide an agentic workflow across four feedback loops: testing, conversion, benchmarking, and optimization. In a production deployment at a major consumer technology company, rewriting a CPU UDF across a 600-billion-row workload reduced operator time from 274.55 hours to 12.78 hours—a 22x speedup affecting over 3,000 jobs. The Spark UDF Bottleneck Apache Spark's declarative SQL and DataFrame APIs allow the NVIDIA cuDF plugin to parallelize row operations on GPUs transparently. However, opaque user-defined functions (UDFs) cannot be optimized by query planners, forcing data to convert from columnar to row-wise format and execute row-by-row on the CPU. This creates round-trip transfer overhead and forfeits GPU acceleration. Three Paths to GPU-Accelerated Rewrites Developers can eliminate fallbacks through three primary rewrite approaches based on logic complexity. Simple operations can be replaced with direct Spark SQL expressions to stay within Spark's expression tree on the GPU. Moderately complex logic can implement the RapidsUDF interface using cuDF's evaluateColumnar method and Java API. Highly complex logic can dispatch to custom CUDA kernels via the Java Native Interface (JNI). AI Agent Skills and Production Results Vectorized APIs and manual verification often make rewriting UDFs labor-intensive. The cuDF skills for Apache Spark provide project templates, test harnesses, benchmark infrastructure, and best practices to guide AI agents across testing, conversion, benchmarking, and optimization loops. A major consumer technology company used this framework to rewrite more than 30 production UDFs, achieving a 22x speedup on a 600-billion-row job affecting over 3,000 workloads.

利用AI智能体改写Apache Spark UDF以实现GPU加速 NVIDIA介绍了利用AI智能体配合cuDF技能将Apache Spark自定义函数（UDF）改写为GPU加速版本的方法，可解决传统UDF因回退CPU执行而造成的性能瓶颈，在实际生产工作负载中实现超过20倍的速度提升。 传统Spark UDF会破坏声明式抽象，迫使GPU管道将列式数据转为行式并回退到CPU执行，造成传输与计算损耗。 UDF改写包含三种路径：简单的SQL表达式替换、通过RapidsUDF接口调用cuDF列式API，以及使用JNI调用底层CUDA内核。 NVIDIA推出cuDF skills for Apache Spark，通过四个具备验证驱动的反馈循环指导AI智能体完成脚手架搭建、转换、基准测试及优化。 某大型消费科技公司应用该流程改写了30多个生产UDF；在一项涉及6000亿行、影响3000多个作业的工作负载中，算子运行时间从274.55小时缩短至12.78小时，加速达22倍。 Spark UDF在GPU加速中的性能瓶颈 NVIDIA cuDF插件利用GPU的SIMD并行能力加速Spark SQL与DataFrame查询。然而，常规UDF包含不可预测的用户代码，破坏了声明式抽象。在GPU加速管道中，系统不得不将列式批处理转换为行式数据并回退至CPU逐行执行，带来往返数据传输开销并丧失了GPU算力加速的机会。 GPU加速UDF的三种改写途径 改写UDF可根据复杂度选择不同方案：简单关系逻辑可直接重写为Spark SQL表达式，以便cuDF将其内联替换为GPU实现；中等复杂度（如字符串解析）可实现RapidsUDF接口，在evaluateColumnar方法中组合cuDF的Java列式API；对于极复杂的逻辑，则可通过JNI调用自定义CUDA内核以获得极致控制。 AI智能体驱动的改写流程与实际成效 向量化列式改写通常耗时耗力，NVIDIA为此提供cuDF技能工具，将工作流分解为测试、转换、基准测试和优化四个阶段的闭环，同时提供项目脚手架与GPU编程实践指导。某消费科技公司利用智能体在数小时内完成了30多个生产UDF的改写，在处理6000亿行数据的场景下实现了22倍的加速，显著降低了运维与时间成本。

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. 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. 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. 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. The UDF challenge 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, 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: Figure 1.A CPU UDF forces an otherwise GPU-resident pipeline to convert columnar batches to rows, execute on the CPU, and convert back to columnar format Thus in an otherwise accelerated query, UDFs have the following drawbacks: The UDF incurs a round-trip data transfer overhead due to fallback. The UDF forgoes potential compute speedup, had it been written using GPU-accelerated primitives. Accelerating UDFs 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 . We outline three approaches to rewrite UDFs: Approach Complexity Best For SQL Low, highest level of abstraction Simple relational logic, e.g. small or inlined functions cuDF Medium, more flexibility and control over execution Most cases, e.g. moderately complex logic like string parsing, nested types CUDA High, maximum flexibility and kernel-level control Most complex or performance critical logic Table 1. Comparison of SQL, cuDF, and CUDA approaches for rewriting Spark UDFs, organized by complexity and best-fit use case 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 """) Rewriting the UDF with 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 """) 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. Rewriting the UDF with cuDF 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 } } 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. Rewriting the UDF with CUDA 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 . Efficient rewrites with AI agents 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. 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 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: Figure 2. Four skills guide agents through testing, conversion, benchmarking, and optimization, with verification-driven feedback loops at each stage Scaffolding and design 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. 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 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: Figure 3. Rewriting a production CPU-based UDF with cuDF reduced operator time from 274.55 to 12.78 hours, a 22x speedup across a 600-billion-row workload 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 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. 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 ‘*’ 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 . 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 .