Implement Arrow PyCapsule Interface #5070

kylebarron · 2023-11-13T01:16:41Z

Which issue does this PR close?

Closes #5067.

Rationale for this change

The PyCapsule interface provides a much safer interface for exposing memory to/from Python, and additionally will work with any arrow library in Python instead of only pyarrow.

What changes are included in this PR?

Implement Send on FFI_ArrowSchema. Send is required to be able to create a new PyCapsule with an FFI_ArrowSchema.
Conversion from DataType, Field, Schema, Array, RecordBatch, RecordBatchReader, and Table through the PyCapsule interface.

It's not possible (I don't think) to implement exporting to the PyCapsule interface because no Python classes are exposed.

Are there any user-facing changes?

None

tustvold · 2023-11-13T12:05:17Z

arrow-data/src/ffi.rs

+/// # Safety
+///
+/// Must be passed a valid [FFI_ArrowArray].
+pub unsafe extern "C" fn release_array(array: *mut FFI_ArrowArray) {


I feel fairly strongly we should not expose this method publicly, it is very hard to use correctly

Agreed this doesn't seem useful.

The motivation for exporting this was to enable creating new PyCapsules from rust with a destructor that releases the array, but it looks like it should be possible to manually call drop without exporting the underlying function:

fn pycapsule_array_destructor(ffi_array: FFI_ArrowArray, _capsule_context: *mut c_void) { drop(ffi_array) } let ffi_array = FFI_ArrowArray::new(...); let array_capsule_name = CString::new("arrow_array").unwrap(); Python::with_gil(|py| { let capsule = PyCapsule::new_with_destructor( py, ffi_array, Some(array_capsule_name), pycapsule_array_destructor, ) .unwrap(); });

FWIW drop is actually a no-op. If you look at the definition it is just an empty function body that consumes the argument.

You should be able to just use PyCapsule::new and it will work correctly

I'm not familiar with the pycapsule internals, but one thing I wonder is whether it runs any destructor if you use PyCapsule::new instead of PyCapsule::new_with_destructor. I.e. do you need even an empty

fn pycapsule_array_destructor(ffi_array: FFI_ArrowArray, _capsule_context: *mut c_void) {}

to get the callback to take ownership again of the ffi_array to get rust to drop it?

I'm not familiar with the pycapsule internals, but one thing I wonder is whether it runs any destructor if you use PyCapsule::new instead of PyCapsule::new_with_destructor.

Yes, it seems it does.
https://github.com/PyO3/pyo3/blob/c8fdb806300ff82747f58bcb9ca08abd3da480b4/src/types/capsule.rs#L273-L282
and
https://github.com/PyO3/pyo3/blob/c8fdb806300ff82747f58bcb9ca08abd3da480b4/src/types/capsule.rs#L78-L84

tustvold · 2023-11-13T12:06:03Z

arrow-schema/src/ffi.rs

+unsafe impl Send for FFI_ArrowSchema {}
+unsafe impl Sync for FFI_ArrowSchema {}


Are we confident these are actually safe to be added?

No, I can't say I understand Send/Sync well enough to be 100% confident this is safe.

The motivation is that right now it's impossible to create a PyCapsule instance from an FFI_ArrowSchema because it's not Send. (See PyCapsule::new).

let ffi_schema = FFI_ArrowSchema::try_from(&*field).unwrap(); let schema_capsule_name = CString::new("arrow_schema").unwrap(); Python::with_gil(|py| { let capsule = PyCapsule::new( py, ffi_schema, Some(schema_capsule_name), ) .unwrap(); });

gives

error[E0277]: `*const i8` cannot be sent between threads safely --> src/array/point.rs:35:17 | 33 | let capsule = PyCapsule::new( | -------------- required by a bound introduced by this call 34 | py, 35 | ffi_schema, | ^^^^^^^^^^ `*const i8` cannot be sent between threads safely | = help: within `arrow::ffi::FFI_ArrowSchema`, the trait `std::marker::Send` is not implemented for `*const i8` note: required because it appears within the type `FFI_ArrowSchema` --> /Users/kyle/.cargo/git/checkouts/arrow-rs-53e12341924f0a1c/b1c43a4/arrow-schema/src/ffi.rs:71:12 | 71 | pub struct FFI_ArrowSchema { | ^^^^^^^^^^^^^^^ note: required by a bound in `pyo3::types::PyCapsule::new` --> /Users/kyle/.cargo/registry/src/index.crates.io-6f17d22bba15001f/pyo3-0.20.0/src/types/capsule.rs:78:29 | 78 | pub fn new<T: 'static + Send + AssertNotZeroSized>( | ^^^^ required by this bound in `PyCapsule::new`

That code is incorrect, IIUC. ffi_schema is a stack-allocated Rust FFI_ArrowSchema. You're trying to store its pointer into a PyCapsule with Python-driven lifetime. It's a good thing that Rust is preventing it :-)

Are we able to just add Send. I'm more confident that is true than Sync.

Hmm, scratch that. PyO3's new_with_destructor actual wraps all its arguments using Box::new, so lifetime should be ok...

Whilst it is never explicitly stated in the C Data interface, I think it is a fairly safe assumption that data can be sent between threads - i.e. it isn't relying on thread-local state or non-atomic reference counts, etc... Happy to be corrected on this front though.

I'm more confident that is true than Sync.

Sync is definitely a stronger assumption that we don't appear to need to make here, although I also struggle to devise a sane example where it wouldn't hold.

tustvold · 2023-11-13T12:07:34Z

Perhaps we could get some of the people who designed PyCapsule to weigh in here, from my reading of it this appears to primarily be an API for python libraries to communicate with each other? I also am very confused by what the ownership semantics of this API are, does the PyCapsule remain the owner? This feels like it can't traverse an FFI interface safely, so I am very confused as to how this is supposed to work...

Perhaps @jorisvandenbossche or @pitrou ??

arrow-data/src/ffi.rs

pitrou

I'm not a Rust programmer, but LGTM at a high level.

arrow/src/pyarrow.rs

arrow-data/src/ffi.rs

pitrou · 2023-11-13T12:14:04Z

arrow-data/src/ffi.rs

+/// # Safety
+///
+/// Must be passed a valid [FFI_ArrowArray].
+pub unsafe extern "C" fn release_array(array: *mut FFI_ArrowArray) {


Agreed this doesn't seem useful.

pitrou · 2023-11-13T12:15:07Z

@WJones @milesgranger You may want to look at this.

pitrou · 2023-11-13T12:38:40Z

Perhaps we could get some of the people who designed PyCapsule to weigh in here, from my reading of it this appears to primarily be an API for python libraries to communicate with each other?

It is designed as a Python-specific protocol (in the Python sense: it uses well-known special methods) around the C Data Interface, and automates conversion between producers and consumers of Arrow data. There is no requirement that the producer or consumer is implemented in Python.

The two main points of this protocol are:

the protocol does not depend on any particular implementation of Arrow; the consumer can consume any PyCapsule provider in an implementation-agnostic manner (it does not need to write any PyArrow-specific code, for example)
the PyCapsule has a destructor that releases the underlying C Data Interface structure if it wasn't moved already; this makes this protocol safer than passing raw pointers as integers (if an exception occurs between exporting and importing, for example, a leak is avoided)

In practice:

provider P exposes a Python object that has a __arrow_c_schema__ method
consumer C is invoked with said Python object, calls __arrow_c_schema__ and imports the C Data Interface structure pointed to by the returned PyCapsule; consumer C does know anything about provider P

I also am very confused by what the ownership semantics of this API are, does the PyCapsule remain the owner?

The PyCapsule just contains a raw pointer. The semantics of the pointer embedded in a PyCapsule are usage-dependent. In this case, the ownership semantics are those of the C Data Interface itself: the contents can be moved by setting the release pointer to NULL.

tustvold · 2023-11-13T12:41:51Z

The PyCapsule just contains a raw pointer. The semantics of the pointer embedded in a PyCapsule are usage-dependent. In this case, the ownership semantics are those of the C Data Interface itself: the contents can be moved by setting the release pointer to NULL.

Ok, so there is an implicit contract that a PyCapsule is never used multiple times?

pitrou · 2023-11-13T12:43:15Z

No, the explicit contract is that a C Data Interface struct can only be used once :-)
https://arrow.apache.org/docs/format/CDataInterface.html#released-structure

The PyCapsule is just a Python-safe way of transporting that struct. It does not change anything to the core semantics.

tustvold · 2023-11-13T12:47:28Z

It does not change anything to the core semantics.

Ok this is what I needed to know, PyCapsule has the same release behaviour as the C Data interface, the copy and custom release behaviour led me to believe it was something reusable.

arrow/src/pyarrow.rs

pitrou · 2023-11-13T12:57:29Z

Yes, "copy" is a misnomer, that's why I suggested "move" instead as a method name.

pitrou · 2023-11-13T13:03:25Z

Should there be some unit tests for this? You can trivially create objects exposing this protocol using PyArrow as a backend, for example:
https://github.com/apache/arrow/blob/aa9f8ecdaf3cb8eb2fb478abff32fb96eb350305/python/pyarrow/tests/test_types.py#L1231-L1239
https://github.com/apache/arrow/blob/aa9f8ecdaf3cb8eb2fb478abff32fb96eb350305/python/pyarrow/tests/test_table.py#L585-L590

tustvold · 2023-11-13T13:05:12Z

Should there be some unit tests for this

The arrow-pyarrow-integration-testing contains some testing infrastructure for this sort of test

Co-authored-by: Raphael Taylor-Davies <[email protected]>

arrow/src/pyarrow.rs

arrow-data/src/ffi.rs

wjones127 · 2023-11-13T16:53:42Z

arrow-schema/src/ffi.rs

+unsafe impl Send for FFI_ArrowSchema {}
+unsafe impl Sync for FFI_ArrowSchema {}


Are we able to just add Send. I'm more confident that is true than Sync.

wjones127 · 2023-11-13T16:55:25Z

arrow/src/pyarrow.rs

Potential refactor for later, but it would be nice to move the Capsule imports to their own functions that aren't named as if to by PyArrow-specific. Right now this module is named for PyArrow, but the import logic will apply to other Arrow implementations as well.

Co-authored-by: Will Jones <[email protected]>

…into kyle/arrow-ffi-copy

pitrou · 2023-11-13T17:06:06Z

Unfortunately, looking at the PyO3 PyCapsule-wrapping APIs, this probably can't work, because PyO3 creates its own boxed structure for the capsule pointer:
https://docs.rs/pyo3/latest/src/pyo3/types/capsule.rs.html#115-119

Non-Rust implementations will expect the capsule pointer to point to a ArrowSchema or ArrowArray struct, but they will get a pointer to some arbitrary Rust structure instead.

This seems to make PyO3's PyCapsule wrapper generally impractical for communicating with non-Rust runtimes.

Edit: scratch that, CapsuleContents starts with the referred T value. Hmm...

tustvold · 2023-11-13T17:06:44Z

arrow/src/pyarrow.rs

@@ -262,9 +263,10 @@ impl FromPyArrow for ArrayData {
            validate_pycapsule(array_capsule, "arrow_array")?;

            let schema_ptr = unsafe { schema_capsule.reference::<FFI_ArrowSchema>() };
-            let array_ptr = unsafe { array_capsule.reference::<FFI_ArrowArray>() };
+            let array_ptr = array_capsule.pointer() as *mut FFI_ArrowArray;
+            let array_mut = unsafe { array_ptr.as_mut() };


FWIW you could use https://doc.rust-lang.org/std/ptr/fn.replace.html as we're already playing with pointers here

Updated in e7ed58d; is that what you meant?

kylebarron · 2023-11-14T16:28:33Z

Thanks! I'll try to get to integration tests tonight

kylebarron · 2023-11-15T04:19:42Z

I added some tests; do those look ok? Do you want to test against old versions of pyarrow as well? We're effectively no longer testing _export_to_c and _import_from_c because the rust code will always prefer the pycapsule methods if they exist. Do you want to add an entry to the matrix to test pyarrow <14 as well?

I had to bump the python version in CI to test with pyarrow 14 at all because pyarrow 14 requires python >=3.8.

pitrou

Tests LGTM.

@kylebarron Do you plan to expose the reverse operation, for example a ToPyCapsule trait that would allow exporting Rust data to Python without having PyArrow installed?

tustvold

Looks good to me.

I think it would be good to add some test coverage for the pre arrow-14 logic, if that isn't too much effort, but happy for this to be done as a follow up PR, just let me know your preference

arrow/src/pyarrow.rs

tustvold · 2023-11-15T11:44:36Z

arrow/src/pyarrow.rs

+        // Newer versions of PyArrow as well as other libraries with Arrow data implement this
+        // method, so prefer it over _export_to_c.
+        // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
+        if value.hasattr("__arrow_c_array__")? {


I think this will allow passing a StructArray to a method accepting PyArrowType<RecordBatch> provided the StructArray isn't nullable. I don't think this is a problem, RecordBatch is a bit of an oddity, but just an observation

Yeah... I'm not sure if there's a standard way in Arrow to discern between a StructArray intended to be an array, vs a StructArray intended to represent a RecordBatch

kylebarron · 2023-11-15T16:14:59Z

Do you plan to expose the reverse operation, for example a ToPyCapsule trait that would allow exporting Rust data to Python without having PyArrow installed?

Unless there are some changes from the second part of apache/arrow#38010, to support this arrow-rs would need to add new Python classes to represent these objects and include the dunder methods, right? Is there appetite to do this, or do people feel that adding new classes would expand the scope of arrow-rs too much?

For my own work in geoarrow-rs I'll be creating my own python classes from rust to represent different types of geometry arrays (as we're discussing in geoarrow/geoarrow-python#38). For the time being I need my own version of the pycapsule interface to ensure I'm not dropping the extension type (ref #4472), but contributing code upstream here ensures it's valid 😅 (and originally I thought I needed upstream changes).

Co-authored-by: Raphael Taylor-Davies <[email protected]>

kylebarron · 2023-11-15T16:26:20Z

I think it would be good to add some test coverage for the pre arrow-14 logic, if that isn't too much effort

I added pyarrow 13 and 14 to the test matrix. So in pyarrow 13, we should see capsule-based tests be skipped but other tests pass.

You might have a better suggestion for structuring that matrix, as it won't upgrade to pyarrow 15 when that's released?

tustvold · 2023-11-15T16:27:17Z

Is there appetite to do this, or do people feel that adding new classes would expand the scope of arrow-rs too much?

From my understanding these classes would be little more than basic shims, in which case I don't see an issue, but I could be missing something

kylebarron · 2023-11-15T16:33:15Z

From my understanding these classes would be little more than basic shims, in which case I don't see an issue, but I could be missing something

At first glance, I thought it would open a can of worms, but maybe not... Would we only have three shims? One per pycapsule type? So exporting a data type/field/schema would all be exported as an ArrowSchema class; exporting Array/RecordBatch would be exported as ArrowArray and Table/RecordBatchReader would be exported as ArrowStream?

One thing I wonder about: if someone sees an ArrowSchema class, how do they know whether to pass it to the pyarrow.data_type, pyarrow.field, or pyarrow.schema constructor?

tustvold · 2023-11-15T16:47:31Z

Three capsules makes sense to me to match tbe corresponding method names. SchemaCapsule, ArrayCapsule and StreamCapsule perhaps. I think it might even be possible to keep them crate private 🤔

if someone sees an ArrowSchema class, how do they know whether to pass it to the pyarrow.data_type, pyarrow.field, or pyarrow.schema constructor?

Is this form of runtime reflection common, do we need to support it, or is the correct usage implicit from context?

.github/workflows/integration.yml

jorisvandenbossche · 2023-11-15T16:54:38Z

Do you need the shims in arrow-rs? It might be enough to expose a function that creates and returns the capsules, and then another library that depends on arrow-rs (like geoarrow-rs) could use that to get the capsules which they can expose on their objects.

Co-authored-by: Will Jones <[email protected]>

wjones127 · 2023-11-15T18:20:47Z

Let's continue the discussion of exporting to PyCapsules here: #5079

pitrou · 2023-11-15T18:40:22Z

arrow/src/pyarrow.rs

+                PyTryInto::try_into(value.getattr("__arrow_c_schema__")?.call0()?)?;
+            validate_pycapsule(capsule, "arrow_schema")?;
+
+            let schema_ptr = unsafe { capsule.reference::<FFI_ArrowSchema>() };


By the way, I'm curious... does it actually drop the FFI_ArrowSchema so as to call it's release pointer?

Does the pycapsule not handle this, as we don't move the data out of it?

Aha, yes, that's true.

pitrou · 2023-11-15T18:43:00Z

Unless there are some changes from the second part of apache/arrow#38010, to support this arrow-rs would need to add new Python classes to represent these objects and include the dunder methods, right?

You don't need new Python classes, you would just return Python objects for the given capsules.

kylebarron added 3 commits November 12, 2023 19:53

arrow ffi array copy

b1c43a4

remove copy_ffi_array

c61d270

docstring

6521cba

github-actions bot added the arrow Changes to the arrow crate label Nov 13, 2023

kylebarron added 2 commits November 12, 2023 22:24

wip: pycapsule support

92b070a

return

6460701

kylebarron changed the title ~~Add a copy method to FFI_ArrowArray~~ WIP: Implement Arrow PyCapsule Interface Nov 13, 2023

tustvold reviewed Nov 13, 2023

View reviewed changes

arrow-data/src/ffi.rs Outdated Show resolved Hide resolved

pitrou reviewed Nov 13, 2023

View reviewed changes

tustvold reviewed Nov 13, 2023

View reviewed changes

arrow/src/pyarrow.rs Outdated Show resolved Hide resolved

Update arrow/src/pyarrow.rs

dfdcfae

Co-authored-by: Raphael Taylor-Davies <[email protected]>

wjones127 reviewed Nov 13, 2023

View reviewed changes

kylebarron and others added 5 commits November 13, 2023 12:00

remove sync impl

5a4f738

Update arrow/src/pyarrow.rs

8a1a05e

Co-authored-by: Will Jones <[email protected]>

Remove copy()

e109c1a

Merge branch 'kyle/arrow-ffi-copy' of github.com:kylebarron/arrow-rs …

3b95ebc

…into kyle/arrow-ffi-copy

Need &mut FFI_ArrowArray for std::mem::replace

05ea67d

tustvold reviewed Nov 13, 2023

View reviewed changes

kylebarron added 2 commits November 14, 2023 11:26

Fix name

becda12

fix ci

2f7767b

kylebarron added 4 commits November 14, 2023 23:08

Add tests

1e7bcd3

Add table test

ae909fb

skip if pre pyarrow 14

6f01c91

bump python version in CI to use pyarrow 14

f183057

kylebarron changed the title ~~WIP: Implement Arrow PyCapsule Interface~~ Implement Arrow PyCapsule Interface Nov 15, 2023

Add record batch test

107acef

pitrou approved these changes Nov 15, 2023

View reviewed changes

tustvold approved these changes Nov 15, 2023

View reviewed changes

kylebarron and others added 2 commits November 15, 2023 11:22

Update arrow/src/pyarrow.rs

6c44e01

Co-authored-by: Raphael Taylor-Davies <[email protected]>

run on pyarrow 13 and 14

6020247

wjones127 approved these changes Nov 15, 2023

View reviewed changes

.github/workflows/integration.yml Show resolved Hide resolved

Update .github/workflows/integration.yml

2e42926

Co-authored-by: Will Jones <[email protected]>

wjones127 merged commit aff86e7 into apache:master Nov 15, 2023
1 check passed

wjones127 mentioned this pull request Nov 15, 2023

Provide functions to produce Arrow PyCapsules #5079

Open

pitrou reviewed Nov 15, 2023

View reviewed changes

tustvold mentioned this pull request Jan 5, 2024

Support new Arrow PyCapsule Interface for Python FFI #5067

Closed

tustvold mentioned this pull request Jan 31, 2024

RecordBatch conversion from pyarrow loses Schema's metadata #5354

Closed

aljazerzen mentioned this pull request Feb 6, 2024

feat: lutra python bindings PRQL/prql#4174

Merged

		unsafe impl Send for FFI_ArrowSchema {}
		unsafe impl Sync for FFI_ArrowSchema {}

Implement Arrow PyCapsule Interface #5070

Implement Arrow PyCapsule Interface #5070

Conversation

kylebarron commented Nov 13, 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 Nov 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kylebarron Nov 13, 2023 • edited Loading

Choose a reason for hiding this comment

tustvold Nov 13, 2023 • edited Loading

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 Nov 13, 2023 • edited Loading

pitrou left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pitrou commented Nov 13, 2023

pitrou commented Nov 13, 2023

tustvold commented Nov 13, 2023

pitrou commented Nov 13, 2023

tustvold commented Nov 13, 2023

pitrou commented Nov 13, 2023

pitrou commented Nov 13, 2023

tustvold commented Nov 13, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pitrou commented Nov 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kylebarron commented Nov 14, 2023

kylebarron commented Nov 15, 2023 • edited Loading

pitrou left a comment

Choose a reason for hiding this comment

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

kylebarron commented Nov 15, 2023 • edited Loading

kylebarron commented Nov 15, 2023

tustvold commented Nov 15, 2023

kylebarron commented Nov 15, 2023 • edited Loading

tustvold commented Nov 15, 2023 • edited Loading

jorisvandenbossche commented Nov 15, 2023

wjones127 commented Nov 15, 2023

Choose a reason for hiding this comment

tustvold Nov 15, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pitrou commented Nov 15, 2023 • edited Loading

kylebarron commented Nov 13, 2023 •

edited

Loading

tustvold Nov 13, 2023 •

edited

Loading

kylebarron Nov 13, 2023 •

edited

Loading

tustvold Nov 13, 2023 •

edited

Loading

tustvold commented Nov 13, 2023 •

edited

Loading

pitrou commented Nov 13, 2023 •

edited

Loading

kylebarron commented Nov 15, 2023 •

edited

Loading

kylebarron commented Nov 15, 2023 •

edited

Loading

kylebarron commented Nov 15, 2023 •

edited

Loading

tustvold commented Nov 15, 2023 •

edited

Loading

tustvold Nov 15, 2023 •

edited

Loading

pitrou commented Nov 15, 2023 •

edited

Loading