Rust AI & Data Meetup: Shehab Amin, Rebuilding the Spark Lakehouse ecosystem in Rust
Recording: Rust AI & Data Meetup: Shehab Amin, Rebuilding the Spark Lakehouse ecosystem in Rust
So, our core thesis is rebuild the engine, don't just accelerate it. We want to unify data and AI workloads. The JVM makes that not possible due to you know, the JVM having bad interoperability with Python. So, Python generally executes out of process. So, a lot of people accelerate Spark. Well, what's the problem with that? The JVM runtime, right? Startup, garbage collection, heap tuning, out of process Python. So, you're doing serialization, deserialization on every record batch, and you have bolted-on control. You have a really tight and, you know, thin line between the JVM and the accelerator language
So, what does Sale change? We have one Rust native runtime for data and AI. We use Arrow and data fusion data end to end, and we have Python and resource control control inside the engine. So, why does Spark hit the ceiling with AI? Spark was built to process data, right? Not run AI where the data lives. So, the legacy runtime blocks AI. AI then has to leave the engine. So, if you're doing feature engineering or embeddings or, you know, inference or whatever, that has to leave the engine. And then the workaround becomes the architecture. You end up stitching together disparate, you know, workloads for data processing and AI, and you create a Frankenstein-like architecture
So, Sale's a unified distributed multi-model lakehouse. We unify batch processing, stream processing, you know, compute-intensive AI workloads, and ad hoc workloads as well. So, what is Sale, right? It's a drop-in replacement of Apache Spark. You just switch your endpoint, and you keep your code, and we have two deployment modes. So, single node, if you have smaller data workloads or you want to work locally, and cluster mode, if you want you know, want to do or petabyte-scale jobs. So, why do we rebuild the engine instead of bolting native operators onto Spark? Well, if you want to accelerate the JVM, right, you intercept the physical plan, you offload operators to native code, and you have a really strict boundary. Uh you know, for example, Photon is column-based, but Spark's row-based, right? So, you end up having a column or a row tax, and you know, uh by virtue of the architecture, Photon doesn't support UDFs. Uh you're JVM trapped, right? It's already on every batch out-of-process Python execution
You can't escape that. Still, we replace the JVM, so 100% Rust, JVM free, uh you know, end-to-end Arrow and DataFusion, columnar throughout, no row-to-columnar tax, and we have native UDF UDF support in process inside the engine. So, how did we rebuild Spark? Well, we use the Spark Connect protocol. Uh you know, a Spark every Spark client uses Spark Connect. So, uh the Spark Connect client sends up a logical plan, we convert that into an internal spec, then we convert that into DataFusion logical plan, and we use uh you know, Arrow. The stack under the hood, Spark Connect API, custom SQL parser, uh which gives us full Spark SQL coverage using parser combinators, Apache DataFusion, which we extend heavily, uh and actually every DataFusion pre-release is tested on Sale, and then PyO3 for embedded CPython, as well as Apache Arrow and Tokio. So, if you want to take a look at Python UDFs uh on Spark versus Sale, you know, you have two wildly different execution models. With Spark, you have to serialize the batch in the JVM, shipped uh to a separate Python worker out of process, deserialize into Python objects, execute the UDF, reserialize back, ship back, and deserialize in the JVM
Not only is that extremely slow, it's also very limiting, you know, right? JVM doesn't have low-level access to, you know, um uh you know, to the hardware, and so you're you're limited in what you can do. Sale, on the other hand, you know, the engine and CPython share one address space. So, the engine reaches the UDF node, you invoke Python in line, and you return zero copy and zero serialization. So, you know, everyone here knows this, but why does Python decide AI performance? Right, all of the workloads run in Python. Uh you know, vectorization, you know, tokenization, vectorization, model calls, feature engineering, and multimodal and physical AI. And so that that uh execution path being really performant and you know, having native low-level access is really important. So, can we, you know, make Python the bridge between data engineering and AI rather than the bottleneck? So, now we're going to go into sales distributed processing architecture. Uh we use the actor model with single-threaded event loop
Uh and you have workers which manage the workers uh worker state. Uh you know, our workers are stateless, so it's it's really only session-based. And you have a job scheduler, and you have a task assigner. Um you know, uh out of the box we do pipeline shuffle, which is streaming shuffle, and the task runner runs tasks with N slots and slot groups. So, what is the actor model and why does it win? The actor model follows the ethos of the Go programming language. Don't communicate by sharing memory, but share memory by communicating. So, you have single-threaded event loops, uh which allows you to have no locks and no contention. And we built this from scratch by hand in Rust because no actor model libraries existed that we were happy with
And uh due to the design of our architecture, we actually have stateless and replaceable workers. So, looking at a job uh job graph, you have a DAG. You know, the job is the top-level unit. This uh a job is broken down to stages, which uh are then broken down to tasks. Tasks process a single partition of data, and each task may be uh retried and reattempted. Uh task region. So, a task region is a set of tasks across stages, and it's really the unit for scheduling and failure recovery. You can do pipeline shuffle, which is streaming shuffle and doesn't write anything to disk, or you can do blocking shuffle and we do our blocking shuffle object storage based
So, you can do it, you know, any object storage compliant URL will work. So, task slot resource management, we follow the architecture of Flink really closely for our task slots and resource management. So, the concepts such as slot sharing, slot sharing groups will sound very familiar if you've worked with Apache Flink before. And then distributed query planning, I'm going to skip this cuz of time. So, shuffle and sale, we have you know, one architecture but two shuffle modes. Uh First is pipeline shuffle. So, you partition streams between stages in memory, you have one task region and nothing's persisted. So, this is the fastest path, right? Arrow batches flow from worker to worker through steam through streams and everything's done all at once
You know, the whole region's live at once and it's coupled failure failure. So, if one task fails, in flight data's gone and the whole region has to retry together. And then we have blocking shuffle. So, we follow Ray's Exo shuffle paper for blocking shuffle, so it's object storage based and you know, it's still a shuffle exchange, although we get rid of the or we get around the small file right problem by doing a a push based let's see, we do doing a push based architecture and so, you're able to, you know, every shuffle reduces able to collect the streams for its channel from all writers and merges them into one output. And then you persist to object storage. So, you know, the stream manager writes it out as compressed arrow and the new region reads it back. So, reducers connect via forward and blocking and are scheduled after the upstream region completes. So, why how do we get around the small file problem? I went ahead of myself earlier
So, we use push based shuffle economics and object storage with no shuffle service. Uh you know, the it comes from the magnet idea. So, it's the same effect as push-based shuffle uh merging the magnet paper, which also which is also discussed in external shuffle. And again, it's it's expressed entirely with the existing job graph primitives, and we're able to provide highly performant blocking shuffle without the need for a shuffle service. So, what is object storage shuffle value, right? Why use object storage shuffle? You know, persistence is not overhead. It's it's resiliency. So, if you want retries without cascades, you know, failed task reads input from object storage, and object storage can be you know, can be external, or it can be you know, uh attached to Kubernetes as a file system, or it can even be local memory or a local file system. Uh and it allows you to you know, persist it right right adds right overhead, but uh under the preemption of of you know, uh memory pressure for example, it could in theory cut total job run time
So, you have fewer retried stages, smaller clusters, and which means a lot cheaper, and spot-friendly scheduling. So, when should you choose uh pipeline versus blocking shuffle? Uh you want to use pipeline shuffle if the job fits the cluster comfortably, if you can fit everything in memory, workers are stable, and retries are rare, or retries are cheap, and latency matters. So, if you if latency's really important, you're doing a real-time workload or ad hoc workload, you're going to want to use pipeline shuffle. Reach for blocking when uh you you know, workers get preempted, or you have a capacity problem, right? Memory's tight, the job is wider than the cluster, or you have long jobs where cascading retries are really expensive. Uh so, I'm going to skip through most of Lakehouse tables, but we rebuilt Delta Lake and Iceberg entirely in Rust. Uh today we support all of you know, all of reading and writing in Delta Lake and Iceberg, writes in Delta Lake and Iceberg, conditional deletes in Delta Lake, merge statements in Delta Lake, and with updates to come soon. Uh the reason why we rebuilt Delta Lake and Iceberg in Rust is, you know, we realized that lakehouse formats are engine contracts and they require really deep integrations with the engine. You know, metadata is data at the end of the day and scalable processing is required
Uh and and so that's you know, that's why we did that. Great. So, uh if you've ever worked with Postgres or MySQL, you will uh you might be familiar with the system catalog, right? The system catalog allows you to have observability through SQL. It tells you how your engine's running. That kind of concept is not available in Spark, so we you know, we added it to Sale, where you can you know, instead of having a a Spark History Server or having uh you know, a Spark UI, right? We do it through OpenTelemetry, so we have via the system catalog. And so we have information about sessions, jobs, stages, tasks, and workers, and integrates with OpenTelemetry, right? So, it fits very nicely with our decoupled storage and compute uh thesis and our you know, stateless workers um you know, architecture. So, how does it work? It's you know, it's just a table schema with with a table and it's just executed as a physical plan. And we don't have to go in into that cuz of time
Benchmarks. Everyone's favorite thing. So, here you're looking at 2026 ClickBench on uh c6g.4xlarge and the two Spark engines that ran c6g.4xlarge were Comet and Spark. Uh combined run, which is a hot and cold, Sale was 6 and 1/2 times faster. On hot run, Sale was 14 times faster. And uh Sale versus Databricks. So, c6g.4xlarge uh versus Databricks's best hot run on ClickBench. So, Databricks's best hot run was 32 XL, which was uh you know, 256 VCPUs and almost 2 terabytes of memory and sale was almost twice as fast on the hot run with you know, significantly less resources
The reason why we report hot run numbers instead of cold run is because cold run is dominated by data load time which reflects data loading and set up and not query engine performance. So we you know, hot run isolates steady state execution. Derived TPCH 2024, we got to do a 2026 one because sale has gone significantly faster since 2024. The sale is up to eight times faster, zero shuffle spill and at peak use 22 gigabytes of memory. There's the individual query breakdown time. And here's the resource utilization, my favorite slide. So you know, Spark used a steady 55 gigabytes of memory, wrote over 110 gigabytes of disk, sale used at peak 22 gigabytes of memory and that peak memory usage only lasted 1 second and this is zero disk writing the entire time. And uh So stop tuning the JVM and replace it
Thank you everyone. >> [applause]