NVIDIA Shows How AI Agents Can Rewrite Spark UDFs for GPU Speedups

Why UDFs hold back GPU-accelerated Spark

The NVIDIA cuDF plugin for Apache Spark uses the parallelism and high-bandwidth memory of GPUs to speed up data analytics. According to NVIDIA, it works largely because Spark's SQL and DataFrame interfaces are declarative: they separate what a query should produce from how the system executes it. That lets the plugin parallelize row-wise operations on the GPU without changes to user code.

UDFs break this model. They are custom routines containing arbitrary user code, which the system cannot reason about or parallelize in any general sense. When a query hits a UDF, the GPU accelerator has to convert data from a columnar to a row-wise format and fall back to the CPU to run the function row by row, then convert the results back.

NVIDIA lists two drawbacks for an otherwise accelerated query:

• The fallback adds a round-trip data transfer overhead.
• The UDF gives up the compute speedup it could have had if written with GPU-accelerated primitives.

Three ways to rewrite a UDF

NVIDIA outlines three rewrite approaches, which it says can eliminate the fallback overhead and yield speedups of over 20x. They differ in complexity and flexibility.

• SQL: lowest complexity and highest level of abstraction. Best for simple relational logic, such as small or inlined functions. The UDF is replaced by an equivalent Spark SQL expression, which becomes visible in Spark's expression tree so the plugin can swap in GPU implementations during query planning.
• cuDF: medium complexity, with more control over execution. NVIDIA recommends it for most cases, such as string parsing or nested types. The UDF implements the RapidsUDF interface, a contract from the cuDF plugin that requires a single method, evaluateColumnar. The plugin passes a GPU-resident ColumnVector to it, and the function composes operations from the Java API for cuDF.
• CUDA: high complexity and maximum flexibility. Intended for the most complex or performance-critical logic that cannot be efficiently composed from existing cuDF operations. A RapidsUDF can invoke a custom CUDA kernel through the Java Native Interface (JNI).

The post illustrates the first two with a simple Scala UDF that converts Celsius to Fahrenheit. The SQL version replaces the function call with the expression celsius * 9.0 / 5.0 + 32.0, while the cuDF version multiplies and adds on GPU columns using scalar values.

Handing the rewrite to coding agents

NVIDIA acknowledges that rewrites still demand significant engineering effort. They require careful API selection, semantic validation and performance measurement, and columnar APIs can be unfamiliar to practitioners with less experience in vectorized programming. The company argues that much of this work suits coding agents: with proper guardrails, an agent can search unfamiliar APIs, implement candidate rewrites and iterate quickly, while executable tests and performance measurements decide whether each change is acceptable.

The cuDF skills for Apache Spark split the rewrite into four focused feedback loops. Each skill defines both the work to perform and the evidence required before moving to the next stage. The skills cover testing, conversion, benchmarking and optimization.

They also supply a complete project template with the cuDF- and Spark-specific scaffolding needed for end-to-end development, including test and benchmark harnesses, encoded routines for output verification and build configuration. NVIDIA says this lets the agent concentrate on the UDF instead of repeatedly rebuilding specialized infrastructure. The skills further include best practices for GPU programming with cuDF, such as memory management and configuration, as an advisory reference. The agent remains free to explore different implementations based on context.

Reported results and what it means

NVIDIA says a major consumer technology company converted more than 30 production UDFs using this workflow. In one rewrite, operator time fell from 274.55 to 12.78 hours across a 600-billion-row workload. These figures come from NVIDIA's own account.

For teams running Spark on GPUs, the pitch is that UDFs, long a source of CPU fallback, can be converted with less manual effort and with verification built into each step. NVIDIA points readers to installing the cuDF skills, a Data + AI Summit session on coding AI agents for Spark UDF GPU acceleration, and the cuDF plugin getting-started guide.