Pass list of lists as argument to Rust from Python using PyO3

user_13127660 0
Viewed 630

I'm trying to pass a list of lists from Python to Rust using Py03. The function I'm trying to pass it to has this signature:

pub fn k_nearest_neighbours(k: usize, x: &[[f32; 2]], y: &[[f32; 3]]) -> Vec<Option<f32>>

I'm writing bindings for a pre-existing lib, therefore I cannot alter the original code. My current way of doing things is this:

// This is example code == DOES NOT WORK
#[pyfunction] // make a new function within a new library with pyfunction macro
fn k_nearest_neighbours(k: usize, x: Vec<Vec<f32>>, y: Vec<f32>) -> Vec<Option<f32>> {
    // reformat input where necessary
    let x_slice = x.as_slice();
    // return original lib's function return
    classification::k_nearest_neighbours(k, x_slice, y)
}

The x.as_slice() function almost does what I need it to, it gives me a slice of vectors &[Vec<f32>], instead of a slice of slices &[[f32; 3]].

I want to be able to run this Python code:

from rust_code import k_nearest_neighbours as knn  # this is the Rust function compiled with PyO3

X = [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [0.06, 7.0]]
train = [
        [0.0, 0.0, 0.0],
        [0.5, 0.5, 0.0],
        [3.0, 3.0, 1.0],
        [4.0, 3.0, 1.0],
    ]

k = 2
y_true = [0, 1, 1, 1]
y_test = knn(k, X, train)
assert(y_true == y_test)
pyo3 rust python
1 Answers

Taking a look at the signature of k_nearest_neighbours shows that it expects [f32; 2] and [f32; 3] which are arrays, not slices (e.g. &[f32].

Arrays have a statically known size at compile time, whereas slices are dynamically sized. The same holds for vectors, you can't control what length the inner vectors have in your example. Therefore you end up with a fallible conversion from the input vectors to the expected arrays.

You can convert from a slice to an array using TryFrom, i.e.:

use std::convert::TryFrom;
fn main() {
    let x = vec![vec![3.5, 3.4, 3.6]];
    let x: Result<Vec<[f32; 3]>, _> = x.into_iter().map(TryFrom::try_from).collect::<Result<Vec<_>, _>>();
}

Putting this together, your function will need to return an error on incorrect input and you will need to create a new vector with arrays that you can pass to your function:

#[pyfunction] // make a new function within a new library with pyfunction macro
fn k_nearest_neighbours(k: usize, x: Vec<Vec<f32>>, y: Vec<f32>) -> PyResult<Vec<Option<f32>>> {
    let x = x.into_iter().map(TryFrom::try_from).collect::<Result<Vec<_>, _>>();
    let y = y.into_iter().map(TryFrom::try_from).collect::<Result<Vec<_>, _>>();
    // Error handling is missing here, you'll need to look into PyO3's documentation for that
    ...
    // return original lib's function return
    Ok(classification::k_nearest_neighbours(k, &x, &y))
}
user_106930220
Related