GH-38325: [Python] Implement PyCapsule interface for Device data in PyArrow

[jorisvandenbossche]

Rationale for this change

PyArrow implementation for the specification additions being proposed in #40708

What changes are included in this PR?

New __arrow_c_device_array__ method to pyarrow.Array and pyarrow.RecordBatch, and support in the pyarrow.array(..), pyarrow.record_batch(..) and pyarrow.table(..) functions to consume objects that have those methods.

Are these changes tested?

Yes (for CPU only for now, #40385 is a prerequisite to test this for CUDA)

• GitHub Issue: [Python] Expose the device interface through the Arrow PyCapsule protocol #38325

[jorisvandenbossche]

This part is a bit repetitive with the non-device version. I could factor that out into a shared helper function

[paleolimbot]

I had a look through here and it is looking great! I don't see anything out of place (but I am also only a little bit familiar with the existing code). I don't personally mind the repeated cast bit (as long as the repetition is tested, which it looks like it is).

Just for this PR I prototyped something similar in nanoarrow ( apache/arrow-nanoarrow#409 ), and the only difference I see is that the device_id for the CPU is -1. I can't find in the spec exactly what it should be...is -1 the accepted value?

import pyarrow as pa
# ! pip install "https://github.com/paleolimbot/arrow-nanoarrow/archive/414bbc44d3e84ecac2807713438d6988ff4d5245.zip#egg=nanoarrow&subdirectory=python"
import nanoarrow as na
from nanoarrow import device

# Wrapper to prevent c_device_array() from falling back on __arrow_c_array__()
class DeviceArrayWrapper:
    def __init__(self, obj):
        self.obj = obj

    def __arrow_c_device_array__(self, requested_schema=None):
        return self.obj.__arrow_c_device_array__(requested_schema=requested_schema)
    

pa_array = pa.array([1, 2, 3])

device.c_device_array(DeviceArrayWrapper(pa_array))
#> <nanoarrow.device.CDeviceArray>
#> - device_type: 1
#> - device_id: -1
#> - array: <nanoarrow.c_lib.CArray int64>
#>   - length: 3
#>   - offset: 0
#>   - null_count: 0
#>   - buffers: (0, 2199023452480)
#>   - dictionary: NULL
#>   - children[0]:

[jorisvandenbossche]

Thanks for testing!

the only difference I see is that the device_id for the CPU is -1. I can't find in the spec exactly what it should be...is -1 the accepted value?

That's a good point. In practice this comes from our implementation in C++:

arrow/cpp/src/arrow/device.h

Lines 82 to 86 in 434f872

	/// \brief A device ID to identify this device if there are multiple of this type.
	///
	/// If there is no "device_id" equivalent (such as for the main CPU device on
	/// non-numa systems) returns -1.
	virtual int64_t device_id() const { return -1; }

But we should probably clarify that in the spec whether that's allowed / expected.

I see that DLPack actually specified that as to be 0 for CPU: https://dmlc.github.io/dlpack/latest/c_api.html#c.DLDevice.device_id
Maybe we should follow that; will open a separate issue.

[jorisvandenbossche]

I updated this PR to include the **kwarg handling.

I think in the meantime, we also have stream support for the device interface, so that could be added as well (although given this is already quite big, I can also do that in a separate follow-up PR). I should also still explicitly test this on CUDA.

[paleolimbot]

Is this check easy to do? (If the failure mode is a crash maybe this would be good to do?)

[jorisvandenbossche]

Yes, we actually expose a device_type on Array and RecordBatch, so we can easily validate this and raise an informative error when trying to cast to requested_schema for non-CPU data.

[paleolimbot]

Thanks!

[jorisvandenbossche]

Thanks for the review! Going to merge this then, and open follow-up issues for expanding to the stream interface and CUDA tests

[conbench-apache-arrow]

After merging your PR, Conbench analyzed the 7 benchmarking runs that have been run so far on merge-commit 1815a67.

There was 1 benchmark result indicating a performance regression:

• Commit Run on ec2-m5-4xlarge-us-east-2 at 2024-06-27 00:42:01Z
  • partitioned-dataset-filter (R) with dataset=dataset-taxi-parquet, language=R, query=vignette

The full Conbench report has more details. It also includes information about 42 possible false positives for unstable benchmarks that are known to sometimes produce them.