Calling into Rust from python


#1

Hi guys,

I’m learning Rust by solving a relatively simple problem and coming from c# world.
I’d have a function in Rust that accepts two non-mutable arrays and calculates a third one which is also a return value.
I’d like to call this code from Python.
Note that I don’t know Python as well - it is just a problem that I come across.

I was thinking that Python should pass all three arrays to avoid memory allocations problem (return array should be freed somehow).
Something like:

#[repr(C)]
pub struct Args {
    vector: *const i32,
    mask: *const i32,
    result: *mut i32,
    len: usize
}
#[no_mangle]
pub extern fn arr(value: Args) -> i32 {
}

In Python I’d use CFFI to call into this, no ideas how to define the Args class or cdef the function yet.

Does this approach make sense?

Thanks


Is it possible to call Rust from python script and how?
#2

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`.

#3

On the Rust side, I suggest using [T] instead of Vec<T>. If you accidentally trigger the Vec's reallocation or deallocation, you’ll probably crash the program, since Python and Rust won’t be using the same allocator.


#4

Ah, that makes a lot of sense. To make sure I fully understand in terms of Rust types and terminology: use [T] – one of the fixed size slice (std::slice) views into memory?


#5

Yep. std::slice::from_raw_parts does the same thing as the Vec version, but it makes it clear that your Rust code isn’t allowed to do anything but read the memory.


#6

I haven’t tested the code but it seems that the case here is wrong:

def _as_f64_array(array):
    """ Cast np.float64 array to a pointer to float 64s."""
    return ffi.cast("double*", array)

array is a numpy structure that holds metadata besides the pointer to the actual data, and so the casting should look something like:

ffi.cast('double *', array.ctypes.data)

Right?


#7

On the other hand… ffi.cast may be smart enough to look past the standard PyObject header, and then the data pointer happens to be the very first element of the array structure, so it all works.


#8

You’re absolutely right. My original post there was a transcription error – I actually used:

def _as_f64_array(array):
    """ Cast np.float64 array to a pointer to float 64s."""
    ffi.cast("double*", array.ctypes.data)

as you suggested.


#9

PSA: when crafting C signatures for Rust code, consider generating them using Rust Cheddar.


#10

Very cool post!
I’m a bioinformatician, where a lot of my biologically-oriented colleagues use python+numpy for bioinformatics, and then come complaining (to me :slight_smile: ) that it’s not fast enough when they do some computation on 1e12 elements…
I’m going use your post as a great example how to drop down into lower languages! Thanks!

One thing that worries me: won’t this part go wrong on 32-bit systems?

The way I understood it, usize is the pointer-sized unsigned int, so size_t in C-parlance.

If I read my wikipedia correctly, “long” is “at least” 32 bits, meaning it’s probably exactly 32 bits on most 32-bit architectures… Only “long long” is guaranteed to be at least 64 bits.

How would one write it to use explicit bit-sizes? On the rust-side, It’s as simple as replacing “usize” with “u64”.
The c-side, I guess should be “uint64_t” instead of “unsigned long”?

Of course, this is mostly hypothetical, since anyone doing any serious numbercrunching wants to be on a 64-bit system to have more than 4GB RAM, but I’m betting someone, somewhere will develop on their (32-bit) ancient laptop and have strange bugs that can’t be reproduced on the (64-bit) desktop.


#11

You’ll probably be interested in Rust-Bio.


#12

I am now! The related crate rust-htslib also looks promising.
Thanks for bringing it to my attention!


#13

You may also be interested in this paper/thesis about using the FST data structure to store protein sequences (PDF; beginning in English; rest in Dutch).

There are also bindings from Python to burntsushi’s fst library.


Crate of the Week
#14

You are a veritable font of useful links! Fst looks much exactly like what my colleague needs! And the summary reads promising.
Mapping structures with lots of overlap in keys :smile:
Dutch won’t be a problem, That’s my native tongue!