[Python] Expose the device interface through the Arrow PyCapsule protocol

[jorisvandenbossche]

We added a new protocol exposing the C Data Interface (schema, array and stream) in Python through PyCapsule objects and new dunder methods __arrow_c_schema/array/stream__ (#35531 / #37797).

We recently also expanded the C Data Interface with device capabilities: https://arrow.apache.org/docs/dev/format/CDeviceDataInterface.html (#34972).

The currently merged PyCapsule protocol uses the stable non-device interface, but so the question is how to integrate the device version in the protocol in order to expose the C Device Data Interface in Python as well. Some options:

1. Only support the device versions going forward (like currently only the cpu version is supported, i.e. the returned capsules always contain a device array/stream).
   _{(this is a backwards incompatible change, but given we labeled the protocol as experimental, we can still make such changes if we think this is the best long-term option. The capsule names would reflect this change, thus this will generate a proper python error if a consumer or producer would not yet have been updated, and we can actually first deprecate the non-device support in pyarrow before removing it. All to say that AFAIU this is perfectly possible if we want it.)}

2. Add separate dunder methods __arrow_c_device_array__ and __arrow_c_device_stream__, and then it is up to the producer to implement those dunders if they can (and we can strongly recommend doing that, also for CPU-only libraries), and to consumers to check which ones are present.

3. Allow the consumer to request a device array with some keyword (eg __array_c_array__(device=True)), which gives the consumer the option to request it while also still giving the producer the possibility to raise an error if they don't (yet) support the device version.

4. Support both options in the current methods without keyword, i.e. allow __arrow_c_array__ to return both a "arrow_array" or "arrow_device_array" capsule (and their capsule name distinguishes both). With the recommendation to always return a device version if you can, but allowing producers to still return a cpu version if they don't support the device one. This only gives some flexibility to the producer, and no control to the consumer to request the CPU version (so this essentially expects that all consumers will handle the device version)

Options 2/3/4 are probably just variants of how to expose both interfaces, and thus the main initial question is whether we want to, long term, move towards an ecosystem where everyone uses the C Device Data Interface, or to keep using both interfaces side by side (as the main interchange mechanism, I mean, the device interface of course still embeds the standard struct).

[jorisvandenbossche]

Looking at the DLPack python page (https://dmlc.github.io/dlpack/latest/python_spec.html), they also define a __dlpack_device__ dunder method, and have a stream keyword in the __dlpack__ method. Do we need something similar here for the sync event? My understanding is that we don't need anything here in addition to the actual C struct, since the sync_event is already part of that? (cc @kkraus14)

[pitrou]

I would in favor of the second approach (separate dunder methods). Python methods are cheap, and these are really two different protocols.

[paleolimbot]

I would in favor of the second approach (separate dunder methods). Python methods are cheap, and these are really two different protocols.

I also would favour a separate dunder method. A consumer that is expecting an ArrowDeviceArray or stream almost certainly has the capability to construct one (i.e., the producer not implementing the method is not a barrier).

[kkraus14]

The downside of using a separate dunder method is that you run into ecosystem fragmentation and inefficiencies. I.E. there was a similar effect that happened with __array_interface__ (https://numpy.org/doc/stable/reference/arrays.interface.html) and __cuda_array_interface__ (https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html). Libraries and applications primarily developed against numpy despite the __array_function__ protocol which made them think and develop in a CPU centric way despite just using array APIs. GPU libraries then needed to either implicitly copy to CPU to support __array_interface__ or raise an Exception to try to guide the developer in the right direction, neither of which is a good developer UX.

Compare that to dlpack (https://github.com/dmlc/dlpack) and the associated Python protocol surrounding it (https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html) which has had device support as a primary goal from its inception. Deep Learning libraries that have multiple device backends like PyTorch, Tensorflow, JAX, etc. all support dlpack but they don't support __array_interface__ or __cuda_array_interface__ because of those issues. Similarly, the __dataframe__ protocol (https://data-apis.org/dataframe-protocol/latest/API.html) has device support as a primary goal as well.

[pitrou]

I don't get the fragmentation and inefficiency argument. Implementing the C Data Interface is necessary to also implement the C Device Data Interface (since the latter is based on the former). Almost no implementation effort is saved by eschewing the former to go with the latter.

In any case, the C Data Interface already exists in C++, C#, Go, Java, Rust, nanoarrow and several third-party projects... It's simply logical to expose it in Python as well.