Zero Copy getter for struct of ndarrays in rust to python

[sdrap]

I am wondering if the following is possible (simplified version of my problem)

• I have an external library with structs of ndarrays with many functions.
• I want to expose it to python (initialize, run the functions, get the struct members) Something Polars seems to be doing for arrows.

So far I can do it, but every time it seems to be a full copy back and forth between python and rust.

use ndarray::Array1;
use numpy::{PyArray1, ToPyArray};
use pyo3::prelude::*;

// the original has 15+ different kind of array with different dimensions
pub struct Test {
    pub x: Array1<f64>
}
impl Test {
    pub fn new(x: Array1<f64>) ->Test {
        Test {x}
    }
}


// now the python class
#[pyclass]
pub struct PyTest{
    pub inner: Test
}


#[pymethods]
impl PyTest {
    #[new]
    fn new(py: Python<'_>, x: &PyArray1<f64>) -> PyTest {
        // get the length of the PyArray1
        let x = x.readwrite();
        let test = Test { x: x.to_owned() };
        py.allow_threads(|| PyTest { test: test })
    }
    #[getter]
    pub fn get_x(&self, py: Python<'_>) -> Py<PyArray1<f64>> {
        self.test.x.to_pyarray(py).to_owned()
    }
}

What I am searching for is however for the getter to be zero copy so that I can call the field PyTest.x[t] from python without copying the array every time I search an index value or a slice of x.

Is that possible?

Sorry if the question is not that clear I am very confused with the interaction between python and rust and I am very new with rust.

[adamreichold]

I think if you want to have zero-copy interactions with Python, you need to turn this around, i.e. store a NumPy array and create a ndarray::ArrayView when Rust code needs to interact with it. This is necessary at least to have a place where the dynamic array borrow checking can run, e.g.

#[pyclass(frozen)]
pub struct PyTest {
   x: Py<PyArray1<f64>>,
}

#[pymethods]
impl PyTest {
  #[new]
  fn new(x: Py<PyArray1<f64>>) -> Self {
    Self { x }
  }

  #[getter]
  fn x(&self, py: Python<'_>) -> Py<PyArray1<f64>> {
    self.x.clone_ref(py)
  }
}

fn modify_x(test: &PyTest) {
  let x = test.x.readwrite();
  ...
}

[sdrap]

Does this operation is a zero copy one? I mean x is something length a milion and I want to call x[t] hundred of thousands of index. I just want to know if this getter is just a reference operation. Could not figure out from the documentation. Even the trait documentation from_owned_array is not clear about it. When I implement it *PyArray1::from_owned_array(py, self.test.x) it complains about copy trait not implemented, which I suppose is expecting a clone.

[adamreichold]

I just want to know if this getter is just a reference operation.

Yes, it is. This is not specific to rust-numpy, but rather a PyO3 API: Py<T> is a reference-counted pointer to type T in the Python heap. Hence, cloning a Py<PyArray1<f64>> will not clone the underlying array, but just increment its reference count.

PyArray::from_owned_array is zero-copy, but to be so, it consumes the given ndarray::Array which probably where the about about the argument not being Copy comes from.

[sdrap]

Many thanks for the clear answer :)!

If understand it correctly:

• if a struct is written as a [#pyclass(frozen)] and use PyArray then it works zero copy (as above)
• if a struct from a separate library is entirely written with ndarrays and those structs wrapped into a [#pyclass(frozen)] then it needs PyArray::from_owned_array to get those elements, which, even is zero copy in itself, will consume the underlying array and therefore in the use of a separate standalone library will necessitate a copy.

I was wondering if there is a way to have a standalone library with ndarrays and interfacing their struct and implementations with a simple wrapping #[pyclass] zero copy with pyo3 (at least getter).

I need to dive more into this, I am not very familiar with rust and these subtleties with the gil interaction is truly difficult.

Many thanks in any cases.

[adamreichold]

if a struct is written as a [#pyclass(frozen)] and use PyArray then it works zero copy (as above)

Sorry if I mislead you there but the frozen part is completely optional for this to be zero-copy. It is just more efficient if you only store references to Python objects, i.e. Py<T>. The guide has some information on that.

will consume the underlying array and therefore in the use of a separate standalone library will necessitate a copy.

You mean if using that library prevents you from consuming the arrays?

I was wondering if there is a way to have a standalone library with ndarrays and interfacing their struct and implementations with a simple wrapping #[pyclass] zero copy with pyo3 (at least getter).

It can only work the other way around because while we are able to produce an ndarray::ArrayView that is backed by NumPy array, an ndarray::Array cannot be stored on the Python heap as it is.

PyArray::from_owned_array wraps the ndarray::Array with a special helper object to be able to put it onto the Python heap and have a NumPy array backing it. But it is highly difficult to recover ownership of the special helper object once we handed the wrapping NumPy array to any Python code due to pervasive shared mutability in Python. This is similar to the problems described in #221

[sdrap]

I see, many thanks for the detailed answer, I am starting to understand the difficulty here. I will have to find a smart way to do this or rewrite the library itself.

I am actually asking for the dual part of what apparently PyArray is doing. Namely exposing to python views (or read only arrays) from rust-ndarray. It seems indeed to be quite difficult to manage ownership.

In any cases many thanks, I start to understand the difficulty to manage both languages heap management.