Advise for zero copy between numpy / arrow-rs

[edgarriba]

HI! I'm running kornia-rs a Rust computer vision library in which i recently moved away from the ndarray crate to represent our internal Tensor and Image structures towards supporting the arrow-rs project.

I have implemented the following approach in order to provide minimal python support to interact with numpy. And I'm looking for adivse wether this could be most efficient / elegant way to achieve zero-copy conversions. Thanks in advance !

pub type PyImage = Py<PyArray3<u8>>;
//pub type PyImage<'a> = Bound<'a, PyArray3<u8>>;


/// Trait to convert an image to a PyImage (3D numpy array of u8)
pub trait ToPyImage {
    fn to_pyimage(self) -> PyImage;
}


impl<const C: usize> ToPyImage for Image<u8, C> {
    fn to_pyimage(self) -> PyImage {
        Python::with_gil(|py| unsafe {
            let array = PyArray::<u8, _>::new_bound(py, [self.height(), self.width(), C], false);
            // TODO: verify that the data is contiguous, otherwise iterate over the image and copy
            std::ptr::copy_nonoverlapping(self.as_ptr(), array.data(), self.numel());
            array.unbind()
        })
    }
}


/// Trait to convert a PyImage (3D numpy array of u8) to an image
pub trait FromPyImage<const C: usize> {
    fn from_pyimage(image: PyImage) -> Result<Image<u8, C>, ImageError>;
}


impl<const C: usize> FromPyImage<C> for Image<u8, C> {
    fn from_pyimage(image: PyImage) -> Result<Image<u8, C>, ImageError> {
        Python::with_gil(|py| {
            let pyarray = image.bind(py);


            // TODO: we should find a way to avoid copying the data
            // Possible solutions:
            // - Use a custom ndarray wrapper that does not copy the data
            // - Return direectly pyarray and use it in the Rust code
            let data = unsafe {
                match pyarray.as_slice() {
                    Ok(d) => d.to_vec(),
                    Err(_) => return Err(ImageError::ImageDataNotContiguous),
                }
            };


            let size = ImageSize {
                width: pyarray.shape()[1],
                height: pyarray.shape()[0],
            };


            Image::new(size, data)
        })
    }
}

[Ekin-Kahraman]

Both directions in the snippet copy the pixel data: copy_nonoverlapping copies into a NumPy allocation, and to_vec() copies back into a Rust allocation. Avoiding those copies requires preserving the owner of the original allocation, rather than just converting its pointer.

For NumPy → Rust processing during a function call, accept a borrowed array and let the algorithm operate on a view:

use numpy::PyReadonlyArray3;
use pyo3::{exceptions::PyValueError, prelude::*};

#[pyfunction]
fn sum_rgb(image: PyReadonlyArray3<'_, u8>) -> PyResult<u64> {
    let view = image.as_array();
    if view.shape()[2] != 3 {
        return Err(PyValueError::new_err("expected H x W x 3"));
    }
    Ok(view.iter().map(|&x| u64::from(x)).sum())
}

That also handles strided arrays. An algorithm requiring contiguous interleaved pixels can use image.as_slice()?, but must reject or explicitly copy non-contiguous inputs. The current owned Image::new(size, data) interface is a separate constraint: borrowing a NumPy view doesn't automatically make it an owned Kornia image. Keep the view/borrow guard alive for its use, and don't allow concurrent external mutation of the underlying array.

For Rust → NumPy, if you have an owned Vec<u8>, move it through ndarray rather than allocating and copying again:

use numpy::{ndarray::Array3, IntoPyArray, PyArray3};

fn owned_rgb<'py>(
    py: Python<'py>, pixels: Vec<u8>, height: usize, width: usize,
) -> PyResult<Bound<'py, PyArray3<u8>>> {
    let image = Array3::from_shape_vec((height, width, 3), pixels)
        .map_err(|e| PyValueError::new_err(e.to_string()))?;
    Ok(image.into_pyarray(py))
}

This is an ownership-transfer example, not an Arrow-buffer-to-Vec conversion. At the Kornia revision you linked, storage wraps a ScalarBuffer; it cannot simply be reinterpreted as an owned Vec. For an Arrow-backed NumPy view, a Python base object must retain the actual buffer owner. borrow_from_array supports that pattern, but is unsafe: lifetime, no reallocation, and read/write aliasing must be upheld. In particular, shared immutable Arrow storage must not be exposed as a freely writable NumPy array.

I compiled and exercised the examples above with numpy 0.30.0 / PyO3 0.30.1, including pointer equality across the owned transfer, strided inputs, wrong dimensions/dtype, and channel validation. They use the current API; the 2024 Kornia binding would need version-appropriate names. Ownership-transfer implementation.

Prepared and tested with Codex; the complete Kornia/Arrow wrapper was not implemented here.