Sharing a view of mutable data from Rust to Python

[Palmik]

Hello,

Here's my use case: I am building up arrays (MutablePrimitiveArray, but I am open to alternatives) in Rust, item by item, and I need to share the data with python (I call the python code through pyo3). The python code only only needs to read the data.

Using ffi::export_array_to_c requires changing the MutablePrimitiveArray to Arc<dyn Array> via as_arc() or into_arc(), which empties it.

fn array_to_python(py: Python, arrow_module: &PyModule, array: ArrayRef) -> PyResult<PyObject> {
    let array_ptr = Box::new(ffi::ArrowArray::empty());
    let schema_ptr = Box::new(ffi::ArrowSchema::empty());

    let array_ptr = Box::into_raw(array_ptr);
    let schema_ptr = Box::into_raw(schema_ptr);

    unsafe {
        ffi::export_field_to_c(
            &ArrowField::new("", array.data_type().clone(), true),
            schema_ptr,
        );
        ffi::export_array_to_c(array, array_ptr);
    };

    let array = arrow_module.getattr("Array")?.call_method1(
        "_import_from_c",
        (array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
    )?;

    unsafe {
        Box::from_raw(array_ptr);
        Box::from_raw(schema_ptr);
    };

    Ok(array.to_object(py))
}

[jorgecarleitao]

Cool question.

In general we can't share mutable data between Rust and Python safely - if Python mutates it at the same time as Rust mutates it, we have undefined behavior 💥.

You mentioned that Python only reads it. In that case, can't you use into_arc as you described - convert it to Arc<dyn Array> and share that? (i.e. MutablePrimitiveArray -> PrimitiveArray -> Arc<dyn Array>)

Or, would you like to afterwards modify it afterwards?

[Palmik]

Thanks for the super fast reply!

I would like to still modify it afterwards, but only in Rust.

Attaching a full non-working example. In the example, I would like arr to keep growing:

use arrow::array::{MutableArray, MutablePrimitiveArray};
use arrow::{array::ArrayRef, ffi};
use polars::prelude::*;
use pyo3::ffi::Py_uintptr_t;
use pyo3::prelude::*;
use pyo3::types::PyDict;

fn example() -> PyResult<()> {
    let mut arr = MutablePrimitiveArray::<i64>::new();
    for i in 0i64..100 {
        arr.push(Some(i));
        Python::with_gil(|py| -> PyResult<()> {
            let arrow_module = py.import("pyarrow")?;
            let foo: Py<PyAny> = PyModule::from_code(
                py,
                "import pyarrow.compute as pc 
def foo(arr):
  print(len(arr))
  print(pc.sum(arr))",
                "",
                "",
            )?
            .getattr("foo")?
            .into();
            let arr = array_to_python(py, arrow_module, arr.as_arc())?;
            foo.call1(py, (arr,))?;
            Ok(())
        })?;
    }

    Ok(())
}

/// Arrow array to Python.
fn array_to_python(py: Python, arrow_module: &PyModule, array: ArrayRef) -> PyResult<PyObject> {
    let array_ptr = Box::new(ffi::ArrowArray::empty());
    let schema_ptr = Box::new(ffi::ArrowSchema::empty());

    let array_ptr = Box::into_raw(array_ptr);
    let schema_ptr = Box::into_raw(schema_ptr);

    unsafe {
        ffi::export_field_to_c(
            &ArrowField::new("", array.data_type().clone(), true),
            schema_ptr,
        );
        ffi::export_array_to_c(array, array_ptr);
    };

    let array = arrow_module.getattr("Array")?.call_method1(
        "_import_from_c",
        (array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
    )?;

    unsafe {
        Box::from_raw(array_ptr);
        Box::from_raw(schema_ptr);
    };

    Ok(array.to_object(py))
}

[jorgecarleitao]

Got it. Interesting use-case :)

I am trying to improve use-cases like yours via #1042 . Until then, I think that the solution is something as follows (not tested!):

use std::sync::Arc;

use arrow2::array::*;

fn to_python(array: Arc<dyn Array>) {
    // the important thing here is that `array` is dropped. 
    // I do not know how Python works exactly, but I would
    // expect that the variable is `del`ed and thus the ref count
    // on the arc is decreased by one, allowing `into_mut` below to work
    println!("{:?}", array);
}

fn main() {
    let mut mutable = MutablePrimitiveArray::<i64>::new();
    mutable.push(Some(1));

    // Arc the data and the array itself
    let array = mutable.into_arc();

    let to_be_passed = array.clone();
    // note that this is passed by value
    to_python(to_be_passed);

    // cast the array
    let primitive = array
        .as_any()
        .downcast_ref::<PrimitiveArray<i64>>()
        .unwrap();

    // clone the ref
    let primitive = primitive.clone();
    drop(array);

    // finally, get a `MutablePrimitiveArray` from `PrimitiveArray`
    let mut mutable = primitive.into_mut().right().unwrap();

    // and mutate it again
    mutable.push(Some(2));
}