How to Rewrite Apache Spark UDFs for GPU Acceleration with AI Agents

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:

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:

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:

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.