GPU Accelerated Spark Connect

Summary

• GPU-accelerated ETL and ML: Spark Connect 4.0 now supports NVIDIA GPU-accelerated plugins for Spark SQL and MLlib, enabling end-to-end GPU acceleration without code changes.
• Dramatic performance gains: Achieve up to 9x faster ETL and ML operations with 80% cost reduction compared to CPU-only baselines on the Databricks Data and AI platform.
• Seamless plugin architecture: The new Spark Connect plugin interface on the Databricks Data and AI platform simplifies adoption of GPU acceleration for both SQL and ML workloads.

GPU Accelerated Spark Connect

Talk by Erik eordentlich, Sr. Manager, NVIDIA ; Gera Shegalov, Principal Distributed Systems Engineer, NVIDIA
Spark Connect, first included for SQL/DataFrame API in Apache Spark 3.4 and recently extended to MLlib in 4.0, introduced a new way to run Spark applications over a gRPC protocol. This has many benefits, including easier adoption for non-JVM clients, version independence from applications and increased stability and security of the associated Spark clusters. The recent Spark Connect extension for ML also included a plugin interface to configure enhanced server-side implementations of the MLlib algorithms when launching the server. In this talk, we shall demonstrate how this new interface, together with Spark SQL’s existing plugin interface, can be used with NVIDIA GPU-accelerated plugins for ML and SQL to enable no-code change, end-to-end GPU acceleration of Spark ETL and ML applications over Spark Connect, with optimal performance up to 9x at 80% cost reduction compared to CPU baselines.
Here’s more to explore:
Production ready data pipelines for analytics and AI: https://www.databricks.com/solutions/data-engineering
The Big Book of Data Engineering: https://www.databricks.com/resources/ebook/big-book-data-engineering-2nd-edition
See all the product announcements from Data + AI Summit: https://www.databricks.com/events/dataaisummit-2025-announcements

Chapters

• 0:00 Introduction and agenda
• 1:34 GPU acceleration background
• 3:57 Spark Connect overview
• 6:02 SQL query execution workflow
• 9:10 ML execution with Connect
• 12:45 Service provider deployment
• 14:11 Spark 4.0 ML innovations
• 18:08 Benchmark results
• 23:30 Demo and conclusion

FAQs

What is Spark Connect?

Spark Connect is a gRPC protocol introduced in Apache Spark 3.4 for SQL/DataFrame APIs and extended in 4.0 for MLlib, enabling remote Spark execution over the network.

How much faster are GPU-accelerated workloads?

With NVIDIA GPU acceleration through Spark Connect, you can achieve up to 9x performance improvements while reducing compute costs by 80% compared to CPU baselines.

What is the Spark Connect client-server architecture?

The Spark Connect client-server architecture separates compute from client interaction, allowing queries from any language to be executed on a remote Spark cluster. This video explains how the client sends computation plans to the server, which optimizes and executes them efficiently while returning results back to the client. This separation enables lightweight clients, improved resource management, and more flexible deployment options compared to traditional embedded Spark drivers.

What's the status of Spark Connect in Spark 4?

Spark Connect has reached production readiness and is now a core component of the Spark ecosystem in version 4.0 and beyond. This video demonstrates significant performance improvements including GPU acceleration through partnerships with NVIDIA, enabling faster execution of data processing and ML workloads. The maturity of Spark Connect makes it a reliable choice for enterprises building scalable distributed applications on modern hardware.

Full transcript

0:00 Introduction and agenda

Welcome everybody. So. I'm the room manager for this session and today uh we are welcoming and Gira from Nvidia who will talk about how to accelerate Spark Connect by using GPUs. Guys, thanks. Hello good morning. Welcome to our presentation on GPU accelerated spark connect. I am with the Spark ETL team at NVIDIA and my co-presenter is or Dentlick who leads the our sister team for Spark ML at Nvidia. This presentation is the result of the collaboration of our teams and the wider Spark community.
There are usual a couple of formalities. First, an acknowledgement. The presentation may contain forward-looking statements. Next, a quick reminder to provide feedback for every session you attend, including this one. It can be done either through the app or on the web in the attendee portal. Our agenda today includes a brief introduction and background with motivation for GPU acceleration um and the Spark Connect and show how Spark Connect can further simplify access to GPU acceleration in Spark SQL and Spark ML context.
We conclude with presenting some benchmark results and a recorded demo. Feel free to catch us in the hallways to see a demo live on the laptop. One reason we need uh scalable distributed systems like Spark is that data is growing rapidly. IDC global data sphere 2023 says that the amount of data um is growing at a compound annual growth rate of 28% per year. Hundreds of zetabytes of data are generated each year. How is compute going to keep up with this growth?

1:34 GPU acceleration background

We believe the answer is to CPU to GPU accelerate batch computing workloads and specifically Apache Spark. So how do we get there? Spark 3 enabled some key functionalities. The first was columner processing which is much more efficient for data processing than a rowbyrow approach. The second was the ability to schedule task on resources such as CPU, FPGA or other type accelerators like a GPU. A third was a plug-in interface to allow alternative SQL execution engines.
These innovations made Rapids accelerator for Apache Spark possible. It accelerates Spark SQL and dataf frame operations using Nvidia GPUs. It works with various spark schedulers, standalone, yarn, kubernetes and the best part there is no code change required. It is integrated with a number of spark distribution and cloud service providers as you can see on this slide. Uh a small selection of it. So the plugin builds on top of Nvidia open source project called rapids uh lib QDF specifically which is a data frame library built for GPUs and it's a native library in C++ we integrated via JNI.
More recently the Spark accelerator for Spark has been extended to also accelerate Spark's AML lib API. This part of accelerator is a Python library built on top uh of the rapids project called QML. The rapids accelerator for patches spark for both ETL and ML is open source. You can find it on GitHub. The links will be provided at the end of the presentation. Uh and you can download binary artifacts like jars and pip packages from Maven and uh Pippi repositors respectively.
Meanwhile, Spark Connect was introduced in Spark 3. 4 and provided a way to run Spark SQL and dataf frame application over lightweight gRPC protocol. It included necessary protobuff definition and client and server code to for executing this interactions with the Spark JVM runtime on the server side remotely. Spark connect decouples the application code from the runtime with very lightweight client uh or application side libraries and enables easy implementation of new language bindings uh in languages like Go, Swift or recently proposed Rust and makes maintenance of Spark versions and client applications much more manageable.

3:57 Spark Connect overview

A key innovation in the just released Spark 4. 2 is the extension of the Spark Connect protocol to cover Spark AML lip functionality. Applications running over Spark Connect can now with Spark 4. 0 include Spark MLIP invocations in addition to the originally supported SQL and data frame operations. For now, ML over connect is supported only in uh Python client via the familiar pispark. mml API, but other language bindings should be possible.
So the tremendous standalone benefits of GPU accelerated Spark and Spark Connect respectively strongly motivate enabling GPU accelerated Spark SQL NML over Spark Connect. The details of which are the main topic of this talk as shown on the right of this figure. The basic idea is simple and involves replacing the baseline CPUbased Spark operations invoked by the Spark Connect server with the GPU accelerated ones. In the rest of the presentations, we explain how and why this works for both SQL and ML, explain key Spark 4.
0 innovations that enable the ML part of this. We delve into service provider deployment and conclude with some benchmarks and an end toend demo of GPU accelerated Spark Connect for ETL and ML together. So before diving deeper, let's recap the benefits of Spark Connect with GPUs. We decoupled apps like notebooks from heavyweight dependencies such as Spark itself, Rapids and CUDA and retaining the familiar Spark user length API client clients and Spark runtime side no longer have to be upgraded in the lock step.

6:02 SQL query execution workflow

So now let's take a look at uh the life cycle of a query executed over Spark Connect. So the execution of Spark SQL and dataf frame based applications. Um on the on the left is the client uh running some simple SQL and dataf frame application code. Uh the application which is pure Python in this case or a Jupyter notebook first connects to the remote Spark connect server. So the flow goes from the left to right. When executed on the client side, the Spark connect converts each Spark SQL query to an unresolved logical plan which is in turn sent to the server over as a gRPC request.
On the server side, the logical plan is resolved into an analyzed logical plan, then optimized and finally converted to a physical plan. A tiny snippet of the physical plan in this example is shown in the light blue shape. So somehow the optimizer turned this count query over join into a hash aggregate. So what happens when we introduce rapid spark the accelerator is able to plug in on the server side at the physical plan level taking an input that is a CPU plan and transforming it to a GPU accelerated plan.
Note that the CPU executes in a rowbyrow basis represented internally as an RDD of internal row. The rapids accelerator instead executes at the column level which is implemented internally as an RDD of columner batch API. Otherwise the physical plans are equivalent with the exact operators being replaced by GPU equivalents. So this is what's happening under the hood. Everything else is using Spark task scheduling and parallel execution uh is driven as usually the client side application code is exactly the same.
No changes required. So that's why we want to have as little green on the left hand side where the client code lives and as much green on the right hand side where the execution is accelerated. And um I am handing it over to work walk you through the SparkML part and the rest of the presentation. Okay, thank you. Good morning everybody. Um actually just uh rearrange this a little bit. Okay. All right. So, now that we've heard about how um SQL, Spark SQL and its acceleration work over Spark Connect, let's take a look at how things look for the for the ML side of things.

9:10 ML execution with Connect

So, here's um CPUbased Spark ML over Spark Connect that was recently enabled in Spark 4. So on the left is a simple Spark ML application that uh uh reads reads in some data from a parquet file and invokes um the fit method on the K means estimator to do some clustering. the execution of the application code um as in the nonML case computes an unresolved logical plan to represent the data feeding into an ML action. In this case, it's K meets fit and uh the plan and the ML parameters in including the algorithm and the action being invoked are serialized and sent via the gRPC protocol to the server.
The server side then des serializes the client's message messages and resolves the logical plan into a data frame along with the spark mlib classes methods and parameters that were specified in the messages. In this example, the fit method is invoked on the k means estimator class on the resolve data frame and its execution just uh follows the or or is carried out by the uh baseline classical spark mlib implementation.
The resulting model after the fit method completes is cached locally on the server side and an ID identifying this model is returned to the client to allow it to um possibly remotely invoke methods on the model like transform for example or extracting and downloading the model coefficients for um um saving locally on the client. Okay, so let's look at the accelerated case. The in this case the Spark Connect server is now attached to a GPU accelerated Spark cluster similarly to the um accelerated SQL case while everything else stays the same especially on the client side just again just like in the SQL case.
So in other words, no application code change is needed and moreover the messages flowing back and forth between the client and server unchanged as well to leverage the attached GPU accelerated uh spark cluster ML on the server side. there's a plug-in mechanism that can be specified to replace the baseline Spark um uh ML uh functionality and with accelerated counterparts and the materialization of the data frame uh to be used in the ML computations leverages all the accelerated physical plan execution engine by the Spark rapid SQL plugin as for so now that we have a basic understanding of how GPU accelerated Spark Connect works from the application user point of view, let's look at deployment from the service provider point of view.

12:45 Service provider deployment

Oops. So one one way to deploy a Spark Connect server is to use a shell script that's conveniently included in the Spark distribution. So here's an example invocation for starting a CPUbased Spark Connect server connected to a standalone Spark cluster. The master node for which is specified in in the URI in the dash-master argument here. The same script can actually be used in the case of GPU accelerated spark connect but with some additional GPU related configs as shown here.
The new configs in in the GPU case are highlighted in dark green. So first notice that we're specifying two jars for Spark Rapids ML and SQL acceleration respectively along with some config specifying the plug-in classes in these jars to invoke for respectively accelerating the SQL part and the ML part. Finally, we have some GPU resource related configs and these go back to Spark 3 as as mentioned and not shown are possibly additional configs that control the behavior of the uh ML and SQL plugins.

14:11 Spark 4.0 ML innovations

So now that we have an idea of how GPU accelerated uh Spark Connect works and is deployed uh let's take a closer look at the innovations in Spark 4 that uh enabled ML over Spark Connect in the first place and then uh the acceleration thereof. So um basically it made sense with that foundation to talk about acceleration of ML over Spark Connect. So roughly here's the state of spark connect before spark 4. So there was built already after as of spark 3.
4 for there was there was only built-in support for Spark SQL and dataf frames with protocol buffer messages message definitions for conveying SQL plans and commands along with server side code for handling the corresponding requests leveraging core Spark SQL for the actual execution in case it wasn't clear before the partially green Spark SQL block on the server side in this figure denotes the fact that Spark, the Spark Rapids accelerator for SQL could actually interoperate with Spark Connect even before Spark 4 since its implementation is act is at the physical plan level and below and thus kicks in fully on the server side.
Spark for developers then went on to extend the Spark Connect protocol to cover essentially all of Spark MLB functionality. This involved defining new messages to convey ML relations like model transforms as well as ML commands like the fit command um for ML Spark ML estimators. The Spark for PISpark client includes all the functionality needed to serialize deserialize the new messages based on application code and also the server response that comes back after the execution.
On the server side, a new ML handler was added that converts the decoded ML protocol messages into Spark ML library invocations. And also added but not shown here had but was shown earlier is the model cache. Okay. So finally, Spark for developers also added an ML backend plug-in interface to allow the ML handler to instead invoke alternative implementations of Spark ML lib functionality, which in our case leverage GPU acceleration.
Here is the ML plug-in interface definition. As you can see, it's pretty simple with only one method that converts an input MLIB class name to a plug-in supplied alternative. And here's a portion of our specific Spark Rapids GPU accelerated ML plugin. You can see at the top the mapping of baseline Spark ML logistic regression class names to their Rapids accelerated counterparts. The heavy lifting is done in the partially shown definition of these alternative classes that actually implement the plug-in logic.
So in this case, the ones corresponding to logistic regression. The full plug-in implementation has other mappings and alternative algo implementations for the ones we support. The only requirement on the plug-in implementations is that it can be invoked from Scola and Java. In our case, this initial Scola invocation actually hands off the full implementation to a Python C++ library that connects back to the Spark Connect Java process via Py forj all on the server side to carry out the actual computations.
So, let's change gears a bit and recall why we're talking about GPUs in the first place. And it is of course all about the benefits of faster and cheaper processing. And with that I'm going to share some benchmark results for um that are similar to ones that we have shared before um for ETL processing and um and uh ML. So the first benchmark is for Spark SQL dataf frame processing and is called NVIDIA decision support or NDS for short.

18:08 Benchmark results

It's a slight adaptation of the TPCDS data processing benchmark widely used in the industry. However, we stress that we're in no way claiming NDS related results as official TPCDS results as they do not comply with the latter's full specification. You can find the details of the NDS benchmark like how it differs from TPCDS along with code to generate the data for it and run the queries in the GitHub repo shown at the bottom of this slide.
This slide has details about the Spark yarn clusters CPU and GPU that were used to obtain the results I'll show in a bit. These clusters were deployed on EC2 instances depicted in the table. eight instances in both the CPU and GPU cases. Um, the GPU instances each had one Nvidia L4 GPU. We ran the NDS benchmark at the 10K or 10 terabyte scale factor with the data stored on HDFS uh running on these nodes. These are some of the configurations provided here mainly for future reference by folks that might be interested in reproducing the results.
These have been tuned to respectively optimize CPU and GPU benchmark performance. Um this includes configs specific to the rapids accelerator on the right like specifying the plug-in class, a shuffle manager and um sizes of various GPU related buffers. So here are the results. The plot on the left shows running time in minutes for the CPU and GPU clusters in blue and green. Blue for CPU, green for GPU respectively. And the plot on the right shows the respective compute costs.
We see that the GPU cluster was about four and a half times faster than the CPU cluster and about a quarter of the cost as measured by the EC2 instance cost per hour multiplied by the running time. The cost savings comes from the much shorter running time despite somewhat higher per hour costs for the GPU cluster. For the ML speedups, here are some somewhat less formal benchmark results. again similar to ones we've presented previously.
This time we ran our standard suite of ML estimator fit method microbenchmarks which can be found um at in our GitHub repo shown on the slide for the Spark ML algorithms that we currently accelerate. Um, and we did this on a in-house server that had multiple A100 GPUs and we used two of them for these results um, in each each case of each algorithm and this was run over Spark Connect with Spark 4 using the Rapids accelerator for ETL and ML.
The only ETL in this case though is reading parquet uh format um and also converting from array type in the format to vector type. Uh for the CPU side we had 16 cores per GPU. So a total of 32 cores. The blue and light green bars show execution time in seconds for the CPU and GPU cases respectively. And as you can see, there's anywhere from 2x to nearly 70x speed up depending on the ML algorithm. And notice that the y-axis is log scale.
So these are somewhat compressed. Um the variation in speed up depends on how much um compute to IO there is for the individual algorithms and k means in particular is a particularly standout one from this point of view. The algorithm parameters and data dimensions were selected to be fairly high compute regimes to highlight benefits of GPU acceleration. Um let me talk a bit about the dark green bars. So this is performance that would be possible if either some code change was allowed for the GPU acceleration case to or if we extended um the plug-in to not fall back when handling certain uh steps related to vector data type.
So if either of those two um were were addressed or possible, we would be able to achieve the uh faster execution time for these benchmarks shown by the dark green bars. Okay. So with that um we can move on. We're going to move on to a demo. So uh in this demo um we have a notebook that carries out ETL and uh ML in a single notebook end to end uh with and without GPU acceleration over Spark for connect and the notebooks adapted from one that is in one of our team's repos um that has a collection of examples for accelerated spark.

23:30 Demo and conclusion

This particular one involves a data set um that is published by May involving mortgage uh performance and uh the original notebook had uh XG boost for the ML part and we are in this demo we're replacing that by our accelerated logistic regression which is part of our our plug-in. Uh that said, XG Boost can also, as it turns out, be run over Spark Connect and there's actually a talk by some of our Nvidia colleagues later today explaining how that can be done.
But in this demo, we're using logistic regression. So, the demo um was run on the client side was run on a laptop that was connected to a workstation with an RTX A6000 graphics card um with 16 CPU cores for the CPU part. And uh we'll show where that connection happens. Uh these two blocks here briefly show um um kind of sketch out the ETL part and the ML part. The ETL part has five steps that do some feature normalization.
Determine the label which is a multiclass label that specifies the degree to which or um by how many months a mortgage is delinquent and the ML task is to predict that um degree of delinquency from the features. Um and then the third step involves um extracting some information about the origin of the loans. Then everything gets joined and the final step is to convert some string features to numerical indices. The ML part um does a train test split on the output of the ETL part and then runs a pipeline consisting of a sparse feature hasher vector assembler and sparse logistic multiclass regret uh logistic regression.
Uh and then finally the pipeline fit method is invoked triggering that whole cascade of operations. So, at this point, I'm going to switch over to uh this other laptop because it's easier to show the video of our demo. Um, so let's see if this works. Okay. All right. So, um, I've got some Let me just start the video and then I'll bouncing back and forth between that doc. Okay, so here is the notebook. Um, and the video is running.
You can let me just pause it real quick. So, at the top, by the way, feel free to chime in at any point, G. So at the top we've got two flags that uh toggle some behavior of the notebook. So the first one is an accelerate on GPU. This one actually um the only the only thing this does is identify which run is GPU or CPU for purposes of logging the runtime at the end of the notebook. It doesn't actually affect what gets executed in the core part.
And then the second flag um uh optionally turns on or off the conversion from the original data set which is CSV to um uh park K. But we're we're in this demo we're starting from park K. So let me jump to uh just move this over a little bit. Jump to Oops. So I think the I think the flag also activates uh the ML backend classes in the connect session. Okay. Um so here's the first key point in the notebook. At the bottom we see that um the uh the line that actually establishes the connection to the remote uh Spark connect server like we saw in the short snippets earlier in the slide.
Um the next key step is at the uh one minute mark. So we'll fast forward there. um where we are s setting up an SSH tunnel. Uh you can see kind of yeah on the sort of bottom of the u terminal window here uh from the laptop to the uh uh workstation where the Spark Connect server is running. And there's the famous 152 port for the Spark Connect server. And then we also tunnel into the Spark uh UI. Okay. So, let's uh scroll forward to uh this point here where we are starting up the CPU run.
Oh, sorry, the CPU cluster. Correct. Yeah. And uh turned out that we had previous to that there was a GPU cluster a GPUbased spark connect server running. And you can see this um display here at the bottom right which is a program called nvtop that shows GPU memory and uh compute utilization. And you can see that uh just as the CPU base server so GPU previous GPU server is terminated and the CPU one starts you can see the GPU activity going to zero.
Okay. So we'll let that run a little bit and now you can see the logging from the uh Spark CPU Spark connect server and now it's kind of just hanging out there waiting. Okay. So the next um key point I will oops forward is so we've we've executed the cell that connects to the to the server and now we're going to start running cells logic cells in the notebook. So the next uh interesting thing is this smoke test where we're doing some very simple um spark data frame processing just to make sure things are working.
So the key thing to look at here is um I guess we're skip to the clamp. Yeah, that's right. Well, ah okay, here we go. So you can see that little bit of uh of um SQL ran and you can see that the physical plan here shows no sign of GPU operators. So um yeah, that confirms that we're in CPU CPU mode here. Um so let's see the next interesting thing is at 255. Okay. So here um okay yeah so just the notebook starts at park um and then there's a believe a we are taking a look at the data sizes here.
Um so uh we have actually just kind of small snippets of the full data set and here we're looking at the raw data set which was the original CSV at 11 gig. Um and then the uh parquet ver version of that is much smaller with u thanks to all the compression in parquet. Okay. So let's move ahead to where we actually start executing the core ETL part which is right about here. So we should soon see a lot of activity in the logging which indicates um all of the uh you know the uh CPU ETL processing.
Here are the definitions of the methods that carry out the steps that I had summarized earlier. Are we doing on time? Um some minutes. Okay. Left. Um yeah. So let me with that just skip ahead to where the ETL part um is done for CPU. So that would be at uh right about here. Okay. So the ETL is completed for CPU and it took about um 99 seconds. Okay. So uh next interesting thing is when we start the um CPU fit method for the ML part.
So that's right here in in this cell uh 26. You can see it's running. Um, and that's going to run for a while. Um, so we will not watch the whole thing. Instead, we will skip ahead to um a little bit of logging that's interesting to look at. So specifically right here um we see a little bit of um protobuff string representation of the ML command that uh got sent as part of the execution of this fit. You can see that in the upper right logging window.
Okay. So let's uh skip ahead to 803. And then here we see um uh a top of the JVM execution of the CPU part. We'll just let that run for a little bit. Um and at uh Okay. And right about here the ETL part is the ML part for CPU is done. Okay. So let me now move over just jump ahead to um where we kick things off for the GPU run. We're running a little low on time here. So, this is where we're changing the that flag to GPU.
Um, and then we go ahead and kill the CPU server up on the right and start the GPU one. And we'll soon see some activity. Um hopefully. Okay, guess we're still moving. Okay, so I think I will u Yeah, there we go. there's some GPU activity and the Spark Connect server at this point is waiting and we will jump ahead to when the um ETL GPU part has been started. You can see a lot more activity on the GPU. Um okay. And then we'll jump ahead to when the Okay.
So the ML part should start um running soon. Yeah, there we go. We got to the pipeline fit for ML with GPU and you can see activity happening again on the GPU. At some point it's really going to shift in the high gear. Okay, there it goes right there. Python worker and can point out Python worker and QML login. Okay. Um and um yeah, then we can um skip ahead to the um yeah, so we yeah, GPU activity is maximal at this point.
Um and the ML part is done. So we can skip ahead to um where we show the uh plot the acceleration um 4x for ETL through almost three and a half for the ML part and then just another representation. And that's that's it for the for the demo. Uh we we're kind of low on time. Um so let me switch back here. and just move real quick to our conclusion slide. So we have some QR codes here for getting more information and um for the yeah ETL plugin, the ML plugin and this is a placeholder for where we're going to have a pointer to the open-source version of the demo that I had showed which isn't quite ready yet.
And that concludes our talk. Thank you.

Learn more about the Databricks Data and AI platform.