{"slug": "how-to-rewrite-apache-spark-udfs-for-gpu-acceleration-with-ai-agents", "title": "How to Rewrite Apache Spark UDFs for GPU Acceleration with AI Agents", "summary": "NVIDIA published cuDF skills for Apache Spark, a set of agent-driven skills that rewrite user-defined functions (UDFs) into GPU-accelerable equivalents, achieving 22x+ speedups in Apache Spark jobs. The skills target the UDF bottleneck that forces the cuDF plugin to fall back to CPU row-by-row execution, and NVIDIA outlines three rewrite approaches — SQL, cuDF, and CUDA — ranked by complexity and best-fit use case.", "body_md": "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.\n\nMeeting that demand requires greater data throughput without a proportional increase in cost. [NVIDIA cuDF](https://developer.nvidia.com/topics/ai/data-science/cuda-x-data-science-libraries/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](https://github.com/NVIDIA/cudf-spark) brings this acceleration to [Apache Spark](https://spark.apache.org/), the industry-leading platform for big data processing and analytics, to achieve [better TCO performance](https://docs.nvidia.com/spark-rapids/user-guide/latest/overview.html#performance-cost-benefits) in major workloads.\n\n**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.\n\n**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](https://github.com/NVIDIA/cudf-spark/tree/main/skills), a set of skills to quickly break through the UDF bottleneck for end-to-end development, from testing, to benchmarking, through integration.\n\n## The UDF challenge\n\ncuDF 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](https://en.wikipedia.org/wiki/Single_instruction,_multiple_data) capabilities of GPUs, without any change in user code.\n\nApache 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:\n\nThus in an otherwise accelerated query, UDFs have the following drawbacks:\n\n- The UDF incurs a round-trip data transfer overhead due to fallback.\n- The UDF forgoes potential compute speedup, had it been written using GPU-accelerated primitives.\n\n## Accelerating UDFs\n\nWith 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](#h.v9vs1jjof47e).\n\nWe outline three approaches to rewrite UDFs:\n\n| Approach | Complexity | Best For | \n|---|---|---|\n| SQL | Low, highest level of abstraction | Simple relational logic, e.g. small or inlined functions | \n| cuDF | Medium, more flexibility and control over execution | Most cases, e.g. moderately complex logic like string parsing, nested types | \n| CUDA | High, maximum flexibility and kernel-level control | Most complex or performance critical logic | \n\n*Table 1. Comparison of SQL, cuDF, and CUDA approaches for rewriting Spark UDFs, organized by complexity and best-fit use case*\n\nTo make these approaches concrete, consider a simple Scala UDF that converts temperature from Celsius to Fahrenheit:\n\n```\nfinal class ToFahrenheit\n    extends Function1[java.lang.Double, java.lang.Double]\n    with Serializable {\n\n  override def apply(celsius: java.lang.Double): java.lang.Double = {\n    if (celsius == null) null\n    else celsius * 9.0 / 5.0 + 32.0\n  }\n}\n\nspark.udf.register(\"to_fahrenheit\", new ToFahrenheit)\nspark.sql(\"\"\"\n  SELECT\n    *,\n    to_fahrenheit(celsius) AS fahrenheit\n  FROM temperatures\n\"\"\")\n```\n\n### Rewriting the UDF with SQL\n\nThe most direct solution for simple UDFs is to replace the UDF with an equivalent Spark SQL expression:\n\n```\nspark.sql(\"\"\"\n  SELECT\n    *,\n    celsius * 9.0 / 5.0 + 32.0 AS fahrenheit\n  FROM temperatures\n\"\"\")\n```\n\nThis 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.\n\n### Rewriting the UDF with cuDF\n\nWhen 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`:\n\n``` python\nimport ai.rapids.cudf.{ColumnVector, Scalar}\nimport com.nvidia.spark.RapidsUDF\nimport scala.util.Using\n\nfinal class ToFahrenheit\n    extends Function1[java.lang.Double, java.lang.Double]\n    with Serializable\n    with RapidsUDF {\n\n  override def apply(celsius: java.lang.Double): java.lang.Double = {\n    if (celsius == null) null\n    else celsius * 9.0 / 5.0 + 32.0\n  }\n\n  override def evaluateColumnar(\n      numRows: Int,\n      args: ColumnVector*): ColumnVector = {\n    Using.Manager { use =>\n      val scale = use(Scalar.fromDouble(9.0 / 5.0))\n      val offset = use(Scalar.fromDouble(32.0))\n      val scaled = use(args.head.mul(scale))\n      scaled.add(offset)\n    }.get\n  }\n}\n```\n\nThe 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](https://docs.rapids.ai/api/cudf-java/stable/index-all) operations to process the column in parallel. This approach provides a richer API for more complex processing.\n\n### Rewriting the UDF with CUDA\n\nFor 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)](https://www.ibm.com/docs/en/epfz/6.2.0?topic=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](https://github.com/NVIDIA/cudf-spark-examples/tree/main/examples/UDF-Examples/RAPIDS-accelerated-UDFs).\n\n## Efficient rewrites with AI agents\n\nDespite 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.\n\nFortunately, 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.\n\n### Available skills\n\nThe 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:\n\n*Figure 2. Four skills guide agents through testing, conversion, benchmarking, and optimization, with verification-driven feedback loops at each stage*\n\n### Scaffolding and design\n\nReliable 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.\n\nThe 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.\n\n### Performance results\n\nA 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:\n\n*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*\n\nThis 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.\n\n## Get started\n\nConverting 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.**\n\nTo get started, install the [cuDF skills for Apache Spark on GitHub](https://github.com/NVIDIA/cudf-spark/tree/main/skills), either directly or using the `npx` skill installer:\n\n```\nnpx skills add NVIDIA/cudf-spark --skill ‘*’\n```\n\nFor 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](https://www.databricks.com/dataaisummit/session/coding-ai-agents-supercharge-your-spark-udfs-gpu-acceleration).\n\nTo learn more about GPU acceleration for Apache Spark, follow the [cuDF plugin for Apache Spark getting-started guide](https://docs.nvidia.com/spark-rapids/user-guide/latest/getting-started/overview.html) to unlock better performance and TCO across your Spark pipelines. Visit our webpage to learn more about [NVIDIA cuDF: GPU-Accelerated Structured Data Processing](https://developer.nvidia.com/topics/ai/data-science/cuda-x-data-science-libraries/cudf), or you can read the docs for more information on [NVIDIA cuDF plugin for Apache Spark](https://docs.nvidia.com/spark-rapids/user-guide/latest/index.html).", "url": "https://wpnews.pro/news/how-to-rewrite-apache-spark-udfs-for-gpu-acceleration-with-ai-agents", "canonical_source": "https://developer.nvidia.com/blog/rewrite-apache-spark-udfs-gpu-acceleration-ai-agents/", "published_at": "2026-10-13 21:33:21+00:00", "updated_at": "2026-10-15 00:12:48.136336+00:00", "lang": "en", "topics": ["ai-agents", "ai-infrastructure", "developer-tools", "ai-tools"], "entities": ["NVIDIA", "Apache Spark", "cuDF", "NVIDIA cuDF plugin for Apache Spark", "CUDA", "Scala"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/how-to-rewrite-apache-spark-udfs-for-gpu-acceleration-with-ai-agents", "markdown": "https://wpnews.pro/news/how-to-rewrite-apache-spark-udfs-for-gpu-acceleration-with-ai-agents.md", "text": "https://wpnews.pro/news/how-to-rewrite-apache-spark-udfs-for-gpu-acceleration-with-ai-agents.txt", "jsonld": "https://wpnews.pro/news/how-to-rewrite-apache-spark-udfs-for-gpu-acceleration-with-ai-agents.jsonld"}}