Here is how I've approached a similar situation: calling into Rust from Python for faster array-based computation. (Disclaimer -- there definitely may be some issues with this approach! But it's worked for me so far. I put this together from a number of helpful posts online. Any corrections or suggestions are welcome).
For my use case, I do a lot of I/O and pre-processing in Python and using NumPy. I then drop down into Rust for heavy lifting. I create the arrays in Python as NumPy arrays then pass these into Rust through pointers. Rust then performs in-place mutation.
In lib.rs:
#![crate_type = "dylib"]
use std::mem;
...
#[derive(Debug, Clone, Copy)]
pub struct Args {
arg1: usize,
arg2: f64
}
pub unsafe fn construct_array_1d<T>(input: *mut T, num_elems: usize) -> Vec<T> {
// Create a Rust `Vec` associated with a `numpy` array via the raw pointer. This is unsafe in Rust parlance
// because Rust can't verify that the pointer actually points to legitimate memory, so to speak.
// At least check that the pointer is not null.
assert!(!input.is_null());
Vec::from_raw_parts(input, num_elems, num_elems)
}
pub fn do_computation(args: Args, v1: &Vec<f64>, v2: &Vec<f64>, mut v3: &mut Vec<f64>) {
// Do in-place heavy lifting here.
...
}
#[no_mangle]
pub extern fn rust_entry_fn(arg1: usize, arg2: f64, v1_ptr: *mut f64,
v1_num_elems: usize, v2_ptr: *mut f64, v2_num_elems: usize,
v3_ptr: *mut f64, v3_num_elems: usize) {
...
// Build ergonomic struct from input parameters.
let args = Args { arg1: arg1, arg2 : arg2 };
// Build immutable input vectors from raw parts.
let v1 = unsafe { construct_array_1d(v1_ptr, v1_num_elems) };
let v2 = unsafe { construct_array_1d(v2_ptr, v2_num_elems) };
// Build mutable output vector from raw parts -- will mutate inplace in Rust.
let mut v3 = unsafe { construct_array_1d(v3_ptr, v3_num_elems) };
// Call into Rust function that performs computations and performs in-state
// mutation of the output array `v3`.
do_computation(args, &v1, &v2, &mut v3);
// Don't let Rust free these arrays after leaving this function... This is
// an FFI call from Python and the arrays are actually "owned" and
// garbage collected by NumPy/Python.
mem::forget(v1);
mem::forget(v2);
mem::forget(v3);
}
Then on the Python side, in main.py:
import numpy as np
from cffi import FFI
## Helper functions for Python-Rust FFI through `cffi`.
def _as_f64(num):
""" Cast np.float64 for Rust."""
return ffi.cast("double", num)
def _as_f64_array(array):
""" Cast np.float64 array to a pointer to float 64s."""
return ffi.cast("double*", num)
def _as_usize(num):
""" Cast `num` to Rust `usize`."""
return ffi.cast("unsigned long", num)
## Instantiate `cffi` object.
ffi = FFI()
ffi.cdef(
"""
void* rust_entry_fn(unsigned long arg1, double arg2, double* v1_ptr, unsigned long v1_num_elems,
double* v2_ptr, unsigned long v2_num_elems, double* v3_ptr, unsigned long v3_num_elems);
"""
)
## Go get the Rust library.
lib = ffi.dlopen("../target/release/mylib.dylib")
arg1 = _as_usize(calc_arg1(...))
arg2 = _as_f64(calc_arg2(...))
v1 = calc_input_array_v1(...)
v1_ptr = _as_f64_array(v1)
v1_num_elems = _as_usize(v1.size)
v2 = calc_input_array_v2(...)
v2_ptr = _as_f64_array(v2)
v2_num_elems = _as_usize(v2.size)
v3 = np.zeros_like(v1)
v3_ptr = _as_f64_array(v3)
v3_num_elems = _as_usize(v3.size)
lib.rust_entry_fn(arg1, arg2, v1_ptr, v1_num_elems, v2_ptr, v2_num_elems, v3_ptr, v3_num_elems)
... ## Do stuff with `v3`.
) that it's not fast enough when they do some computation on 1e12 elements..