Design strategy for Python library using pyo3 and Rust to boost performance

[attack68]

For context I have a library written in Python designed for Python. Parts of it are now written in Rust (where it has enhanced performance). I would like to transition more to Rust to continue to speed things up. Creating a fully functioning Rust replicable library is, however, only a secondary goal and not the primary.

The challenge I have is that the bottlenecks end up being in a high layer so to get to those I need to port every lower layer and dependent object to Rust first in order to both maintain compatibility and allow calculations to occur in Rust. But this can introduce slowdowns.

Here is a concrete example. Suppose I create some (supposed) immutable data object in Python and time its data fetching usage...

from datetime import datetime as dt

class PyDataProvider:
    def __init__(self):
        # ... calculation
        self.dates = [dt(2000, 1, 1), dt(2001, 1, 1), dt(2002, 1, 1), dt(2003, 1, 1), dt(2004, 1, 1)]
        self.vals = [1., 2., 3., 4., 5.]
        self.val = 2.5

obj = PyDataProvider()

%timeit obj.dates[3]  # 27.5 ns ± 0.659 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
%timeit obj.vals[3]  # 27.5 ns ± 0.659 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
%timeit obj.val  # 20.4 ns ± 0.369 ns per loop (mean ± std. dev. of 7 runs, 100,000,000 loops each)

If I now port these objects to Rust, there is overall degradation of performance of using them due to the TypeConversion cost and memory clone. Unless the ...calculation part of this makes up some speed it is not really a very good atomic addition.

use pyo3::prelude::*;
use chrono::prelude::*;

#[pyclass]
#[derive(Clone, Default, Debug)]
pub struct RsDataProvider {
    #[pyo3(get, set)]
    pub(crate) dates: Vec<NaiveDateTime>,
    #[pyo3(get, set)]
    pub(crate) val: f64,
    #[pyo3(get, set)]
    pub(crate) vals: Vec<f64>,
}

impl RsDataProvider {
    pub fn new() -> Self {
        // ... calculations
        Self {
            dates: vec![ndt(2001, 1, 1), ndt(2001, 1, 2), ndt(2001, 1, 3), ndt(2001, 1, 4), ndt(2001, 1, 5)],
            val: 3.5_f64,
            vals: vec![1_f64, 2_f64, 3_f64, 4_f64, 5_f64],
        }
    }
}

#[pymethods]
impl RsDataProvider {
    #[new]
    fn new_py() -> PyResult<Self> {
        Ok(DataProvider::new())
    }
}

In Python I now achieve the speeds:

obj = RsDataProvider()
%timeit obj.dates[3]  # 50x WORSE:  1.4 μs ± 23.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
%timeit obj.vals[3]  # 5x WORSE: 138 ns ± 3.39 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
%timeit obj.val  # 2x WORSE: 43.3 ns ± 1.58 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)

For simple float types this isn't that bad but the datetime/ NaiveDateTime is much weaker (and these lists are very small).

I am considering providing an interim solution by using cached Python wrappers.

from functools import cached_property

class PyRsDataProvider:
    vals = cached_property(lambda self: self.obj.vals)
    val = cached_property(lambda self: self.obj.val)
    dates = cached_property(lambda self: self.obj.dates)
    
    def __init__(self):
        self.obj = RsDataProvider()

obj = PyRsDataProvider()
%timeit obj.dates[3]  # ONLY 1.5x WORSE:  42.1 ns ± 1.96 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
%timeit obj.vals[3]  # ONLY 1.5x WORSE:  39.4 ns ± 1.83 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
%timeit obj.val  # ONLY 1.5x WORSE:  32.6 ns ± 0.958 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)

This allows reasonable preservation of speed in Python but the struct is directly available in Rust to continue to perform calculations there. Problem is its quite a bit of work and requires very careful structuring.

Is there another design pattern besides this that I am missing, or is my line of thinking broadly on the right track?

[kitizz]

I'm curious if you ever found a nice solution to this issue? Or if the answer essentially is: more rust

[attack68]

I made use of cached_properties extensively. Transitioned more objects to Rust at lower levels and tried to do as little Python/rust type conversion back and forth as possible. But ultimately this is a bottleneck that cant be wished away, although since writing this I belivev Pyo3 speeds have also improved.