Speeding Up Python with Rust and No Prior Knowledge

Python is a great language and if you prioritize prototyping and getting results quickly (as I do with most of my projects) the trade-offs between programmer productivity and processing speed are fine for me. The really time-consuming parts are outsourced to numpy, OpenCV, or other “fast” libraries with compiled code, and the speed of my Python code (or lack thereof) is negligible. Ultimately, my time spent optimizing is limited, and CPU cycles are cheap. Yet, a few times a year I am stuck with an issue where Python is prohibitively slow. Recently I was working on a piece of code computing hachures lines for vector maps and my Python implementation of the algorithm takes several minutes to complete. However, this is a task where I want to work interactively and experiment to find a visually pleasing result. This only works if I can compare input parameters, modify and observe quickly. In other words: it needs to be fast enough to keep me in the loop.

flowlines example image

I am mostly writing Python, I do not enjoy touching C code and would rather avoid it. But there’s a new kid on the block (… for like, … 10 years) and that’s Rust. In combination with PyO3 (“Pythonium-Trioxide”) as a bindings framework, it should be a straightforward drop-in replacement for some relatively simple Python code. That’s at least what I hoped.

A quick summary of someone getting his first Rust program to run as a Python module:

1. Learning Rust:

Compared to Python, Rust the language and its standard library have a considerably steeper learning curve. Reading some basic info before starting and picking up what’s necessary on the go might not work out well. The Rust Book, however, is a good intro to the language and it made sense to read most of the chapters as a preparation.

Be aware: the book tends to err on the side of having a simpler explanation and be a bit lengthy rather than assuming a certain level of computer science knowledge. I think for this kind of book that’s the right choice, but it’s good to know that beforehand and adapt your reading style accordingly.

The Rust compiler does output some excellent error messages for common problems and is a remarkable help while learning the language. Even though Python has improved in this regard over the last few years substantially, rustc output goes the extra mile to point you in the right direction. However, given the complexity of the language in comparison to Python, this is necessary, though.

2. Tooling

Just to get started the necessary tooling is rather minimal:

3. Project Structure:

Splitting the Rust code and the Python bindings into two separate crates (similar to what is recommended here) felt a lot “cleaner”, reduced the mental complexity, and allowed me to compile and test the Rust code without any Python bindings. That was a big plus in speed and “programming ergonomics”.

project
├── project_py
│   ├── src
│   │   └── lib.rs
│   ├── project_py.pyi
│   └── Cargo.toml
├── project_rs
│   ├── src
│   │   ├── lib.rs
│   │   └── main.rs
│   ├── tests
│   │   └── integration_tests.rs
│   └── Cargo.toml
├── .venv
│   └── ...
├── test_python_bindings.py
├── Cargo.toml
└── README.md

Build and run the rust code:
cargo run --package project_rs --bin project_rs

Build and install the Python library (into the local virtual environment):
maturin develop -m project_py/Cargo.toml

Run the Python test code:
python test_python_bindings.py

4. Python Bindings:

PyO3 in combination with maturin did work out well. Ignoring the “first steps” and just going through the user guide is the approach I would recommend. Afterward, the PyO3 examples make more sense.

A downside of the duplication of objects (one set of “clean” rust structs and methods and one set of objects for PyO3 annotations) is the duplication. I found no simple way of avoiding this boilerplate code using Rust’s partial object-oriented features. This might be a rookie issue.

Though there are bindings for NumPy’s ndarrays as well, that did not work out for me as expected.

5. Python Documentation

At the time of writing type hints for mypy/the IDE need to be created manually as .pyi stub files (see PEP 484).

6. Performance:

Computing hachure lines for a full 3 meter map on an Intel Macbook 16 takes 24s with Rust and 317s running the Python version. That’s a 13x speed-up. Gains are not linear though, for a smaller canvas the Rust implementation is only 2-3x faster.

The main performance issue I faced as a Rust beginner is that I did not figure out how to pass OpenCV images (as numpy n-dimensional arrays) from Python to Rust without an additional copy.

I was slightly surprised by how large the gap between develop (no compiler optimization) and release (all optimizations enabled) builds is (about a 10x speedup as well).

7. Wrapping it up:

I was surprised about the quality of both the Rust book and the documentation. Notably, I had the impression that the Rust community is a pleasant group of people, which is not something that can be taken for granted.

For simple problems that can be boiled down to sequential data processing (passing/copying data from Python to Rust, waiting for the Rust code to finish the computation, and returning it), Rust is basically a drop-in replacement.

However, the typical Python user caveat applies: if you get too used to Python and its built-in support for type variable handling, it takes a conscious effort to make yourself use a “strict” language. Don’t underestimate the complexity and learning curve with Rust.


Resources: