image_alt:Moving from one abstraction layer to another.

PostgresML ismoving to Rust for our 2.0release

PostgresML isMoving to Rust for our 2.0Release

<pclass="author">

September 19, 2022

</p>

PostgresML is a fairly young project. We recently released1.0 and now we're considering what we want to accomplish for2.0. In addition to simplifying the workflow for building models, we'd like to address runtime speed, memory consumption and the overall reliability we've seen for machine learning deployments running at scale.

PostgresML is a fairly young project. We recently releasedv1.0 and now we're considering what we want to accomplish forv2.0. In addition to simplifying the workflow for building models, we'd like to address runtime speed, memory consumption and the overall reliability we've seen is needed for machine learning deployments running at scale.

Python is generally touted as fast enough for machine learning, and is the de facto industry standard with tons of popular libraries implementing all the latest and greatest algorithms. Many of these algorithms (torch, tensorflow, xgboost, numpy) have been optimized in C, but not all of them. For example, most of the[linear algorithms](https://github.com/scikit-learn/scikit-learn/tree/main/sklearn/linear_model) inscikit learnareimplemented in pure Python, although theyrely on numpy, which is a convenient optimization. It also usescython in a few performance critical places. This ecosystem has allowed PostgresML to offer a ton of functionality with minimal duplication of effort.

Python is generally touted as fast enough for machine learning and is the de facto industry standard with tons of popular libraries, implementing all the latest and greatest algorithms. Many of these algorithms (Torch, Tensorflow, XGboost, NumPy) have been optimized in C, but not all of them. For example, most of the[linear algorithms](https://github.com/scikit-learn/scikit-learn/tree/main/sklearn/linear_model) inScikit-Learnarewritten in pure Python, although theydo use NumPy which is a convenient optimization. It also usesCython in a few performance critical places. This ecosystem has allowed PostgresML to offer a ton of functionality with minimal duplication of effort.

<figure>

</figure>

To illustrate ourmotivation, we'll create a test set of 10,000 random embeddings with 128 dimensions, andstore them in a table. Our first benchmark will simulate semantic ranking, by computing the dot product against every member of the test set, sorting the results and returning the top match.

To illustrate ourgoal, we created a test set of 10,000 random embeddings with 128 dimensions andstored them in a table. Our first benchmark will simulate semantic ranking, by computing the dot product against every member of the test set, sorting the results and returning the top match.

```sql linenums="1" title="generate_embeddings.sql"

GROUP BY i %10000;

Spoiler alert, idiomatic Rust is about 10x faster than native SQL,theembedded PL/pgSQL scripting language, andPython in this benchmark. Rust comes close to the handoptimized assembly version of the Basic Linear Algebra Subroutinesimplementation for the dot product. Numpy is supposed to provide optimizations in cases like this, but it's actually the worst performer. Data movement from Postgres to PL/Python is pretty good. It's even faster than the pure SQL equivalent, but adding the extra conversion from Python list to Numpy array takes almost as much time as everything else. Machine Learning systems that moverelativelylarge quantities of data around can become dominated by these extraneous operations, rather than the ML algorithms that actually generate value.

Spoiler alert: idiomatic Rust is about 10x faster than native SQL, embedded PL/pgSQL, andpure Python. Rust comes close to the hand-optimized assembly version of the Basic Linear Algebra Subroutines(BLAS) implementation. NumPy is supposed to provide optimizations in cases like this but it's actually the worst performer. Data movement from Postgres to PL/Python is pretty good; it's even faster than the pure SQL equivalent, but adding the extra conversion from Python list to Numpy array takes almost as much time as everything else. Machine Learning systems that move large quantities of data around can become dominated by these extraneous operations rather than the ML algorithms that actually generate value.

<center>

ORDER BY 1

LIMIT 1;

```

=== "Numpy"

=== "NumPy"

```sql linenums="1" title="define_numpy.sql"

CREATE OR REPLACE FUNCTION dot_product_numpy(a FLOAT4[], b FLOAT4[])

RETURNS FLOAT4

LIMIT 1;

```

We're building with the Rust[pgx](https://github.com/tcdi/pgx/tree/master/pgx) crate that makes our development cycle even nicer than the one we use to manage Python. It really streamlines creating an extension in Rust, so all we havetoworry about is writing our functions. It took about an hour to port all of our vector operations to Rust with BLAS support, and another week to port all the "business logic" for maintaining model training and deployment. We've even gained some new capabilities for caching models across connections (independent processes), nowthat we have accesstoPostgres shared memory, without having to worry about Python's GIL and GC. This is the dream of Apache's Arrow project, realizedforour applications, without having to change the world, just our implementations. 🤩 Single-copy end-to-end machine learning, with parallel processing and shared data access.

After learning this, it should cometono surprisethat we're movingtoRustforPostgresML v2.0.

ML isn't just about basic math and a little bit of business logic. It's about all those complicated algorithms beyond linear regression for gradient boosting and deep learning. The good news is that most of these libraries are implemented in C/C++, and just have Python bindings. There are also bindings for Rust ([lightgbm](https://github.com/vaaaaanquish/lightgbm-rs),[xgboost](https://github.com/davechallis/rust-xgboost),[tensorflow](https://github.com/tensorflow/rust),[torch](https://github.com/LaurentMazare/tch-rs)).

Our new extension will be built with the Rust[pgx](https://github.com/tcdi/pgx/tree/master/pgx) crate which makes our development cycle even nicer than the one we use to manage Python. It streamlined creating an extension in Rust, so all we have to worry about is writing our business logic. It took about an hour to port all of our vector operations to Rust with BLAS support, and another week to port all the rest needed for maintaining model training and deployment. Now that we have access to Postgres' shared memory, we've even gained some new capabilities such as caching models across connections (independent processes), without having to worry about Python's GIL and garbage collector. This is the dream of Apache's Arrow project, realized for our applications, without having to change the world, just our implementations: single-copy end-to-end machine learning, with parallel processing and shared data access. 🤩.

ML isn't just about basic math and a little bit of business logic. It's about all those complicated algorithms beyond linear regression for gradient boosting and deep learning. The good news is that most of these libraries are implemented in C/C++ and just have Python bindings. There are also bindings for Rust ([lightgbm](https://github.com/vaaaaanquish/lightgbm-rs),[xgboost](https://github.com/davechallis/rust-xgboost),[tensorflow](https://github.com/tensorflow/rust),[torch](https://github.com/LaurentMazare/tch-rs)).

<figure>

<imgalt="It's all abstraction"src="/blog/images/abstraction.webp" />

</figure>

The results are somewhat staggering. We didn't spend any time intentionally optimizing Rust over Python. Most of the time spent was just trying to get things to compile. 😅 It's hard to believe the difference is this big, but those fringe operations outside of the core machine learning algorithms really do dominate, requiring up to 35x more time in Python during inference. The difference between classification and regression speeds here are related to the dataset size. The scikit learn handwritten image classification dataset effectively has 64 features (pixels) vs the diabetes regression dataset having only 10 features.**The more data we're dealing with, the bigger the improvement we see in Rust**. We're even giving Python some leeway by warming up the runtime on the connection before the test, which typically takes a second or two to interpret all of PostgresML's dependencies. Since Rust is a compiled language, there is no longer a need to warmup the connection.

The results are somewhat staggering. We didn't spend any time intentionally optimizing Rust over Python; most of the time was spent just trying to get things to compile. 😅. It's hard to believe the difference is this significant, but those fringe operations outside of the core machine learning algorithms really do dominate, requiring up to 35x more time in Python during inference. The difference between classification and regression here are related to the dataset size. The Scikit-Learn handwritten image classification dataset effectively has 64 features (pixels) vs. the diabetes regression dataset having only 10 features.

**The more data we're dealing with, the bigger the improvement we see in Rust**. We're even giving Python some leeway by warming up the runtime on the connection before the test, which typically takes a second or two to interpret all of PostgresML's dependencies. Since Rust is a compiled language, there is no longer such a requirement: all the code is warm and optimized.

<center>

</center>

```sql linenums="1" title="train.sql"

SELECTpgml.train(

project_name=>'Handwritten Digit Classifier',

FROMpgml.digits;

The API is identical between versions1.0 and2.0. We take breaking changes seriously and we're not going to break existing deployments just because we're rewriting the whole project. The only reason we're bumping the major version is because we feel like this is a dramatic change, but we intend to preserve a full compatibility layer with models trained on1.0 in Python.This means that to get the full performance benefits, you'll need to retrain models after upgrading.

The API is identical between versionsv1.0 andv2.0. We take breaking changes seriously and we're not going to break existing deployments just because we're rewriting the whole project. The only reason we're bumping the major version is because we feel like this is a dramatic change, but we intend to preserve a full compatibility layer with models trained onv1.0 in Python.However, this does mean that to get the full performance benefits, you'll need to retrain models after upgrading.

Besides backwards compatibility, we're building a Python compatibility layer to guarantee we can preserve the full Python model training APIs,when Rust APIs are not at parity in terms of functionality, quality or performance. We started this journey thinking that the olderalgorithms in scikit learn that are implementedinvanilla Pythonwould be the best candidates for replacement in Rust, but that is only partly true. There are high quality efforts in[linfa](https://github.com/rust-ml/linfa) and[smartcore](https://github.com/smartcorelib/smartcore) that also show 10-30x speedup overscikit, but they still lack some of the deeper functionality like joint regression, some of the more obscure algorithms andhyperparams, and some of the error handling that has been hardened intoscikit with mass adoption.

Besides backwards compatibility, we're building a Python compatibility layer to guarantee we can preserve the full Python model training APIs,since Rust APIs are not at parity in terms of functionality, quality or performance. We started this journey thinking that the oldervanilla Python algorithmsinScikitwould be the best candidates for replacement in Rust, but that is only partly true. There are high quality efforts in[linfa](https://github.com/rust-ml/linfa) and[smartcore](https://github.com/smartcorelib/smartcore) that also show 10-30x speedup overScikit, but they still lack some of the deeper functionality like joint regression, some of the more obscure algorithms andhyperparameters, and some of the error handling that has been hardened intoScikit with mass adoption.

<center>

</center>

Interestingly, the training times for some of the simplest algorithms aremuchworse in the Rust implementation. Until we can guarantee eachalgorithm implementation is an upgrade in every way, we'll continue to use the Python compatibility layer on a case by case basis to avoid any unpleasant surprises.

Interestingly, the training times for some of the simplest algorithms are worse in the Rust implementation. Until we can guarantee eachRust algorithm is an upgrade in every way, we'll continue to use the Python compatibility layer on a case by case basis to avoid any unpleasant surprises.

We believe that[machine learning in Rust](https://www.arewelearningyet.com/) is mature enough to add significant value now, where we'll be using the same underlying C libraries, andthatit's worth contributing to the Rustimplementations further to helpbringthem up to full feature parity.With this goal in mind, we intend to drop our Python compatibility layer completely in 3.0, and only support 2.0 models trained with Rust long term. Part of our 2.0releaseprocesswill include a benchmark suite for the full API we support via all Python libraries, so that we can track our progress toward pure Rust implementationsacross the board.

We believe that[machine learning in Rust](https://www.arewelearningyet.com/) is mature enough to add significant value now. We'll be using the same underlying C/C++ libraries, and it's worth contributing to the RustML ecosystembringit up to full feature parity.Our v2.0release will include a benchmark suite for the full API we support via all Python libraries, so that we can track our progress toward pure Rust implementationsover time.

Many thanks and ❤️ to all those who are supporting this endeavor. We’d love to hear feedback from the broader ML and Engineering community about applications and other real world scenarios to help prioritize our work.We'd also appreciate your supportin the form of[starsonour github](https://github.com/postgresml/postgresml).

Many thanks and ❤️ to all those who are supporting this endeavor. We’d love to hear feedback from the broader ML and Engineering community about applications and other real world scenarios to help prioritize our work.You can show your supportby[starring usonGitHub](https://github.com/postgresml/postgresml).

<center>

<videocontrolsautoplayloopmutedwidth="90%"style="box-shadow:008px#000;">