Learn How to Rewrite Apache Spark User-Defined Functions for GPU Acceleration with AI Agents

NVIDIA has introduced cuDF skills for Apache Spark, a set of AI-agent skills that guides teams through rewriting user-defined functions into GPU-accelerated equivalents within minutes, with reported 22x+ speedups in Apache Spark jobs. The skills are aimed at data engineers and platform teams that run Apache Spark pipelines on GPUs and need to keep custom business logic from forcing jobs back onto CPUs.

Why custom Spark logic can erase GPU acceleration

Enterprise applications have become more demanding, and agentic AI has increased the need for large volumes of enterprise data, frequent complex queries, and work across increasingly large datasets. That raises the bar for data throughput without a proportional increase in cost. NVIDIA cuDF addresses this challenge by using the massive parallelism and high-bandwidth memory of GPUs for data analytics. The cuDF plugin for Apache Spark brings that acceleration to Apache Spark, a leading platform for big data processing and analytics, with the goal of improving total cost of ownership outcomes in major workloads.

The remaining bottleneck is often custom code. Spark’s declarative programming interface separates what a query should produce from how the system should execute it. That separation lets the cuDF plugin transparently parallelize row-wise operations using GPU parallel capabilities without requiring changes to user code.

User-defined functions break that abstraction. They are custom routines used to extend built-in functionality, which makes them useful for bespoke business logic. But they contain arbitrary user code that the system cannot reason about or parallelize in a general way. During execution, the GPU accelerator must convert data from a columnar format to a row-wise format, fall back to the CPU, and execute the function row by row.

That fallback has two main costs. First, it introduces a round-trip data transfer overhead. Second, the function forgoes the compute speedup it could have had if it were written using GPU-accelerated primitives. In an otherwise accelerated query, a single user-defined function can therefore erase much of the benefit of moving the pipeline to GPUs.

Three rewrite paths for Spark pipelines

With a small, targeted code change, many user-defined functions can be rewritten into a GPU-acceleratable form by using the right primitives. Rewriting removes the accelerator fallback overhead and can yield speedups of over 20x, according to NVIDIA.

There are three main approaches, ordered by complexity and control:

• Spark expression language: low complexity and the highest level of abstraction. It is best for simple relational logic, such as small or inlined functions.
• cuDF APIs: medium complexity, with more flexibility and control over execution. It is best for many common cases, including moderately complex logic such as string parsing and nested types.
• CUDA kernels: high complexity, with maximum flexibility and kernel-level control. It is best for the most complex or performance-critical logic.

A simple Scala user-defined function that converts temperature from Celsius to Fahrenheit illustrates the idea.

For simple cases, the most direct path is to replace the user-defined function with an equivalent Spark expression. This makes the logic visible in Spark’s expression tree. During query planning, the cuDF plugin can replace those expressions with corresponding GPU implementations, keeping data and execution on the GPU.

When higher-level expressions are insufficient or undesirable, the function 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. The cuDF plugin recognizes instances of the RapidsUDF interface and passes a GPU-resident ColumnVector to evaluateColumnar. The function then composes operations from the Java API for cuDF operations to process the column in parallel.

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. NVIDIA notes that examples of such conversions can be found in the cuDF examples for Spark user-defined functions.

Agent skills turn rewrites into a repeatable workflow

Despite the performance benefits, rewriting user-defined functions can require significant engineering effort. It demands careful API selection, semantic validation, and performance measurement. Columnar APIs can also be unfamiliar to practitioners with less experience in vectorized programming.

Much of that 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.

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

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 that knowledge into code, the agent can focus on the function itself rather than repeatedly reconstructing specialized infrastructure.

The skills also include best practices for GPU programming with cuDF, including memory management and configuration. These are intended as an advisory reference rather than a rigid checklist. The agent remains free to explore different implementations based on surrounding context and select an approach based on measured correctness and performance.

What the production numbers suggest

A major consumer technology company used this workflow to convert 30+ user-defined functions used in production data pipelines into GPU-accelerated implementations. The rewrites preserved the behavior of the original CPU user-defined functions while removing costly fallbacks from otherwise GPU-accelerated workloads.

In one example, NVIDIA reports that rewriting a production CPU-based user-defined function with cuDF reduced operator time from 274.55 to 12.78 hours, a 22x speedup across a 600-billion-row workload. That particular function was affecting 3,000+ jobs.

The reported result is not just a one-function benchmark. The value for Spark teams is that GPU acceleration is often limited by pockets of custom logic, not by the main pipeline. The agent skills are designed to make the rewrite process more repeatable, so practitioners can iterate on rewrites on the order of hours rather than days or weeks while preserving correctness.

NVIDIA says the cuDF skills for Apache Spark are available on GitHub and can be installed directly or through the npx skill installer. The company also points to a Data + AI Summit session and a cuDF plugin for Apache Spark getting-started guide for teams evaluating GPU acceleration across Spark pipelines.