feat(api!): make ArrowArrayStreamReader Send #4232

wjones127 · 2023-05-17T03:12:15Z

Which issue does this PR close?

Rationale for this change

The underlying implementation is arguably Send, and making it not Send makes it extremely hard to work with. This drops support for Clone, but if someone wants to sacrifice Send for Clone, they can wrap it in an Arc themselves.

What changes are included in this PR?

Alters implementation of ArrowArrayStreamReader to use Box instead of Arc internally.
Refactors PyArrowConvert convert into a few different traits, so that we can differentiate between conversions that borrow vs those that take ownership.
Cleans up some of the ffi_stream code to remove some unnecessary heap allocations.

Are there any user-facing changes?

This is a breaking change to PyArrowConvert.

tustvold · 2023-05-22T07:04:23Z

I'm not sure about the changes to the python interface, we've run into challenges in the past because python can't transfer ownership. I'll have a play later today and see if we can't make it Send without requiring a breaking change

tustvold

I left some comments, I like where this is headed. I think there is one memory safety issue, but otherwise mostly just nits

arrow/src/ffi_stream.rs

tustvold · 2023-05-22T09:03:45Z

arrow/src/ffi_stream.rs

 //!
 //! // consumed and used by something else...
 //!
 //! // import it
-//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr).unwrap() };
+//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() };


Suggested change

//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() };

//! let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap();

I think we could deprecate from_raw

I think it's still relevant. For example, there are cases where we need to implement a C API where we will be passed a pointer to an ArrowArrayStreamReader. For example, this is done in the "bind stream" function in ADBC:

https://github.com/apache/arrow-adbc/blob/75ba2cef9080b19f4d242d0a03c17db304609d85/rust/src/implement/internal.rs#L647-L664

tustvold · 2023-05-22T09:04:36Z

arrow/src/ffi_stream.rs

@@ -231,8 +227,7 @@ impl ExportedArrayStream {
                    let struct_array = StructArray::from(batch);
                    let array = FFI_ArrowArray::new(&struct_array.to_data());

-                    unsafe { std::ptr::copy(addr_of!(array), out, 1) };
-                    std::mem::forget(array);
+                    unsafe { std::ptr::write_unaligned(out, array) };


I think this will cause the array to be released when array goes out of scope, i.e. immediately. I think eventually this will result in a double-free

This is a good question. However, the docs for write_unaligned say that src is guaranteed not to be dropped:

https://doc.rust-lang.org/std/ptr/fn.write_unaligned.html

Internally, it's calling intrinsics::forget(). So I think this is sound.
https://doc.rust-lang.org/src/core/ptr/mod.rs.html#1447

TIL, I had no idea that was the semantic - makes sense

arrow/src/ffi_stream.rs

tustvold · 2023-05-22T09:08:29Z

arrow/src/pyarrow.rs

        let stream = Box::new(FFI_ArrowArrayStream::empty());
        let stream_ptr = Box::into_raw(stream) as *mut FFI_ArrowArrayStream;

-        unsafe { export_reader_into_raw(Box::new(self.clone()), stream_ptr) };


I see why IntoPyArrow is needed, this needs to "consume" the stream

Exactly. It's a little messy, but I did want to convey we are moving ownership, while all the other conversion are either creating entirely new objects (data types) or sharing ownership (arrays).

tustvold · 2023-05-22T09:09:15Z

arrow/src/pyarrow.rs


-impl<'source, T: PyArrowConvert> FromPyObject<'source> for PyArrowType<T> {
+impl<'source, T: FromPyArrow + IntoPyArrow> FromPyObject<'source> for PyArrowType<T> {


Suggested change

impl<'source, T: FromPyArrow + IntoPyArrow> FromPyObject<'source> for PyArrowType<T> {

impl<'source, T: FromPyArrow> FromPyObject<'source> for PyArrowType<T> {

The reason both traits are there is because PyArrowType<T> requires. If I remove one we get

the trait bound `T: FromPyArrow` is not satisfied

tustvold · 2023-05-22T09:09:29Z

arrow/src/pyarrow.rs

    fn extract(value: &'source PyAny) -> PyResult<Self> {
        Ok(Self(T::from_pyarrow(value)?))
    }
 }

-impl<'a, T: PyArrowConvert> IntoPy<PyObject> for PyArrowType<T> {
+impl<T: FromPyArrow + IntoPyArrow> IntoPy<PyObject> for PyArrowType<T> {


Suggested change

impl<T: FromPyArrow + IntoPyArrow> IntoPy<PyObject> for PyArrowType<T> {

impl<T: IntoPyArrow> IntoPy<PyObject> for PyArrowType<T> {

tustvold · 2023-05-22T09:09:46Z

arrow/src/pyarrow.rs

    }
 }

-impl<T: PyArrowConvert> From<T> for PyArrowType<T> {
+impl<T: FromPyArrow + IntoPyArrow> From<T> for PyArrowType<T> {


Suggested change

impl<T: FromPyArrow + IntoPyArrow> From<T> for PyArrowType<T> {

impl<T> From<T> for PyArrowType<T> {

arrow/src/pyarrow.rs

tustvold · 2023-05-24T11:52:56Z

The integration test failure looks legitimate and would suggest something is off with this PR somewhere...

/home/runner/work/_temp/57085026-34c0-45a7-adab-58efe95c1cf7.sh: line 4: 9787 Aborted (core dumped) pytest -v .
tests/test_sql.py::test_record_batch_reader

wjones127 · 2023-05-24T14:55:38Z

Yes, thanks for the reminder. I'll look into this soon.

* feat(api make ArrowArrayStreamReader Send * simplify ptr handling * rename pyarrow traits to conform to guidelines * pr feedback * remove dangling Box::from_raw

feat(api make ArrowArrayStreamReader Send

b223a7e

github-actions bot added the arrow Changes to the arrow crate label May 17, 2023

simplify ptr handling

a9dc6ec

wjones127 force-pushed the GH-4222-c-stream-send branch from 06753ea to c51d589 Compare May 21, 2023 03:27

rename pyarrow traits to conform to guidelines

0f4aa77

wjones127 force-pushed the GH-4222-c-stream-send branch from c51d589 to 0f4aa77 Compare May 21, 2023 03:36

wjones127 marked this pull request as ready for review May 21, 2023 03:54

tustvold reviewed May 22, 2023

View reviewed changes

tustvold added the api-change Changes to the arrow API label May 22, 2023

pr feedback

0a1c18a

remove dangling Box::from_raw

2dee727

tustvold merged commit 3e5b07a into apache:master May 27, 2023

tustvold mentioned this pull request Jun 2, 2023

Make ArrowArrayStreamReader Send #4222

Closed

wjones127 mentioned this pull request Jul 2, 2023

feat: support RecordBatchReader on boxed trait objects #4475

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(api!): make ArrowArrayStreamReader Send #4232

feat(api!): make ArrowArrayStreamReader Send #4232

wjones127 commented May 17, 2023 •

edited

Loading

tustvold commented May 22, 2023

tustvold left a comment

tustvold May 22, 2023

wjones127 May 23, 2023

tustvold May 22, 2023

wjones127 May 23, 2023

tustvold May 23, 2023

tustvold May 22, 2023

wjones127 May 23, 2023

tustvold May 22, 2023

wjones127 May 23, 2023

tustvold May 22, 2023

tustvold May 22, 2023

tustvold commented May 24, 2023 •

edited

Loading

wjones127 commented May 24, 2023

	//! let stream_reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() };
	//! let stream_reader = ArrowArrayStreamReader::try_new(stream).unwrap();


		impl<'source, T: PyArrowConvert> FromPyObject<'source> for PyArrowType<T> {
		impl<'source, T: FromPyArrow + IntoPyArrow> FromPyObject<'source> for PyArrowType<T> {

	impl<'source, T: FromPyArrow + IntoPyArrow> FromPyObject<'source> for PyArrowType<T> {
	impl<'source, T: FromPyArrow> FromPyObject<'source> for PyArrowType<T> {

	impl<T: FromPyArrow + IntoPyArrow> IntoPy<PyObject> for PyArrowType<T> {
	impl<T: IntoPyArrow> IntoPy<PyObject> for PyArrowType<T> {

	impl<T: FromPyArrow + IntoPyArrow> From<T> for PyArrowType<T> {
	impl<T> From<T> for PyArrowType<T> {

feat(api!): make ArrowArrayStreamReader Send #4232

feat(api!): make ArrowArrayStreamReader Send #4232

Conversation

wjones127 commented May 17, 2023 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

tustvold commented May 22, 2023

tustvold left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold commented May 24, 2023 • edited Loading

wjones127 commented May 24, 2023

wjones127 commented May 17, 2023 •

edited

Loading

tustvold commented May 24, 2023 •

edited

Loading