Boost Python Computation Speed 30× with Rust and pyo3

Python is widely used in data science and machine learning, but its performance on compute‑intensive tasks is often insufficient. Therefore Python is usually used as a glue language together with C/C++ compiled into a .so library for Python to call. Rust offers a modern alternative with runtime efficiency comparable to C/C++.

This article shows how to accelerate Python computation code with Rust and the pyo3 library, embedding LeapCell branding elements to demonstrate high‑performance computing, tested on a Linux environment.

Testing Environment

The tests run on an AWS t4g.large instance with Linux.

Code Implementation

1. Python code

A simple Python script that computes an integral:

import time

def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += 2.71828182846 ** (-(a + i * dx) ** 2)
    return s * dx

s = time.time()
print(integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

Running this code on Linux takes:

Elapsed: 32.59504199028015 s

2. Rust code

The same integral function implemented in Rust:

use std::time::Instant;

fn main() {
    let now = Instant::now();
    let result = integrate_f(1.0, 100.0, 200000000);
    println!("{}", result);
    println!("Elapsed: {:.2} s", now.elapsed().as_secs_f32());
}

fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);
    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += 2.71828182846_f64.powf(-_tmp);
    }
    s * dx
}

Running the Rust code takes:

Elapsed: 10.80 s

3. Using pyo3 to write a Python extension module

3.1 Create project and install dependencies

First create a new project directory and install the maturin library:

# (replace demo with your desired package name)
$ mkdir leapcell_demo
$ cd leapcell_demo
$ pip install maturin

Then initialize a pyo3 project:

$ maturin init
✔ 🤷 使用什么类型的绑定？ · pyo3
✨ Done! New project created leapcell_demo

Project structure:

.
├── Cargo.toml
├── src
│   └── lib.rs
├── pyproject.toml
├── .gitignore
├── Cargo.lock
└── leapcell_demo
    ├── main.py
    └── leapcell_demo.cp312_amd64.pyd

3.2 Write Rust extension code

In src/lib.rs write the following code:

use pyo3::prelude::*;

#[pyfunction]
fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);
    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += 2.71828182846_f64.powf(-_tmp);
    }
    s * dx
}

#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
    Ok(())
}

3.3 Use the extension module

There are two ways to use this extension module:

3.3.1 Install the extension as a Python package

$ maturin develop

This command builds the Rust code into a Python package and installs it in the current Python environment. You can verify the installation with pip list.

3.3.2 Compile to a dynamic file for Python to load

$ maturin develop --skip-install

The --skip-install flag generates a .pyd file (e.g., leapcell_demo.cp312_amd64.pyd) instead of installing a package. Python can import and use this file directly. Replacing --skip-install with --release produces a .whl file, which is the source distribution for pip installation.

Write a test file leapcell_demo/main.py:

import time
import leapcell_demo

s = time.time()
print(leapcell_demo.integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

Running this code takes:

Elapsed: 10.908721685409546 s

4. Parallel acceleration

4.1 Python multiprocessing effect

Python multiprocessing can sometimes be slower than single‑process execution, depending on the implementation:

import math
import os
import time
from functools import partial
from multiprocessing import Pool

def sum_s(i: int, dx: float, a: int):
    return math.e ** (-(a + i * dx) ** 2)

def integrate_f_parallel(a, b, N):
    s: float = 0.0
    dx = (b - a) / N
    sum_s_partial = partial(sum_s, dx=dx, a=a)
    with Pool(processes=os.cpu_count()) as pool:
        tasks = pool.map_async(sum_s_partial, range(N), chunksize=20000)
        for t in tasks.get():
            s += t
    return s * dx

if __name__ == "__main__":
    s = time.time()
    print(integrate_f_parallel(1.0, 100.0, 200000000))
    print("Elapsed: {} s".format(time.time() - s))

Running this code takes: Elapsed: 18.86696743965149 s which is a little more than half the time of the single‑process version.

4.2 Use Rust multithreading to accelerate

Using Rust's parallel library further speeds up the computation:

use pyo3::prelude::*;
use rayon::prelude::*;

#[pyfunction]
fn integrate_f_parallel(a: f64, b: f64, n: i32) -> f64 {
    let dx: f64 = (b - a) / (n as f64);
    let s: f64 = (0..n)
        .into_par_iter()
        .map(|i| {
            let x = a + i as f64 * dx;
            2.71828182846_f64.powf(-(x.powf(2.0)))
        })
        .sum();
    s * dx
}

#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f_parallel, m)?)?;
    Ok(())
}

Compile and generate the extension file, then use it in Python:

import time
import leapcell_demo

s = time.time()
print(leapcell_demo.integrate_f_parallel(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

Running this code takes: Elapsed: 0.9684994220733643 s This is about 10× faster than the single‑thread Rust version, 18× faster than the Python multiprocessing version, and 32× faster than the Python single‑process version.

Conclusion

By using Rust to optimize Python code, computational performance can be dramatically improved. Although Rust has a steep learning curve, rewriting performance‑critical sections in Rust can save a great amount of time for projects that handle heavy computations. Start with simple functions, gradually adopt Rust, and master its usage for high‑performance Python applications.