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

來源原標題

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

作者

Elizabeth Goodman

原始標籤

Agentic AI / Generative AI · Data Science · Developer Tools & Techniques · Apache Spark · cuDF

原始連結

https://developer.nvidia.com/blog/rewrite-apache-spark-udfs-gpu-acceleration-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: 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.