feat(python): Add ArrowDeviceArray extension to Python bindings #313

paleolimbot · 2023-10-31T20:39:30Z

This PR adds basic support for wrapping the ArrowDeviceArray in nanoarrow Python. You have to try pretty hard to get anything that isn't a CPU array here, but 99% of this is just to get a repr() and make the device array easier to debug/understand. In some future the CUDA and Metal implementations could live in separate Python packages as well.

import nanoarrow.device
import pyarrow as pa

nanoarrow.device.device_array(
    pa.record_batch([pa.array([1, 2, 3])], ["col"])
)

<nanoarrow.device.DeviceArray>
- device_type: 1
- device_id: 0
- array: <nanoarrow.Array struct>
  - length: 3
  - offset: 0
  - null_count: 0
  - buffers: (0,)
  - dictionary: NULL
  - children[1]:
    'col': <nanoarrow.Array int64>
      - length: 3
      - offset: 0
      - null_count: 0
      - buffers: (0, 5698482274496)
      - dictionary: NULL
      - children[0]:

Most of the code changes here are to add the appropriate reprs for the Array and Schema objects. These should work for both CPU and Device flavours of the C Data interface.

@jorisvandenbossche the reprs should make it easier for your __arrow_c_array__ / __arrow_c_schema__ testing!

codecov-commenter · 2023-10-31T20:41:12Z

Codecov Report

Attention: 2 lines in your changes are missing coverage. Please review.

Comparison is base (2d28306) 86.91% compared to head (966a014) 89.10%.

Files	Patch %	Lines
python/src/nanoarrow/_lib_utils.py	96.49%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #313      +/-   ##
==========================================
+ Coverage   86.91%   89.10%   +2.19%     
==========================================
  Files          70        4      -66     
  Lines       10591      101   -10490     
==========================================
- Hits         9205       90    -9115     
+ Misses       1386       11    -1375

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

paleolimbot · 2023-11-15T14:42:01Z

If there are no objections, I'll merge this tomorrow!

jorisvandenbossche

That's a nice improvement for the repr!

python/src/nanoarrow/_lib_utils.py

jorisvandenbossche · 2023-11-16T13:09:28Z

python/src/nanoarrow/_lib.pyx

@@ -474,20 +481,8 @@ cdef class Array:
        else:
            return None

-    def view(self):


Is there a reason you are making this into a function instead of method? (for Schema it's still a method, so this gets a bit inconsistent)

Yes...the Array as defined here doesn't know if it's a CPU or GPU array, and the "view" does. I don't know that the view() method is a particularly good idea for the Schema, either...a cached property might be a better fit.

Early on you made the suggestion that the Array and ArrayView were a little confusing, name wise. I wonder if renaming the Array to CArray (+ Schema and ArrayStream) would fix some of that. Then ArrayView could be renamed to Array, which would probably make more sense to everybody.

...the CPU or GPUness matters because the ArrayView will let you access buffer content, which will crash for the GPU.

Early on you made the suggestion that the Array and ArrayView were a little confusing, name wise. I wonder if renaming the Array to CArray (+ Schema and ArrayStream) would fix some of that. Then ArrayView could be renamed to Array, which would probably make more sense to everybody.

I was just writing up some thoughts on that! ;) -> #319

jorisvandenbossche · 2023-11-16T13:17:20Z

python/src/nanoarrow/_lib.pyx

+        return <uintptr_t>&self.c_array
+
+cdef class DeviceHolder:
+    """Memory holder for an ArrowDevice


Is this ArrowDevice concept something specific from nanoarrow? (I don't find it in https://arrow.apache.org/docs/dev/format/CDeviceDataInterface.html)

Yes, it's a nanoarrow invention. It contains methods for things like "copy this buffer to another device" but it's not ABI stable, so we'd have to workshop how to register a device from another package such that Device.resolve() does the right thing.

Can you add some docstring then clarifying that? (to the Device class below)

jorisvandenbossche · 2023-11-17T17:52:35Z

python/src/nanoarrow/_lib.pyx

+    cdef ArrowDevice* _ptr
+
+    def __cinit__(self, object base, uintptr_t addr):
+        self._base = base,


Is this base used anywhere?

It's a strong reference to the DeviceHolder (like in Array, Schema, and ArrayStream). It's probably less confusing to just use the Device for this (unlike Array and Schema, there's never any other base than a DeviceHolder).

Right...in making that change I remembered why this is like this: the CPU device is statically allocated in C, so there's no base/DeviceHolder. But other types of device implementations might not (it's not totally solved in the C library yet).

python/src/nanoarrow/_lib.pyx

Co-authored-by: Joris Van den Bossche <[email protected]>

jorisvandenbossche reviewed Nov 16, 2023

View reviewed changes

paleolimbot and others added 21 commits November 17, 2023 13:19

copy device stuff

de25d62

build with ArrowDeviceArray

239e9e0

building + importing

c4a7b09

with test

7c6ba3e

formatting

bf5cefe

with device array

07234b6

fix isort

90fc9dd

move device stuff

6acfaef

maybe fix pre-commit

a519a57

grrr

f5a296b

maybe fix again

792aa2a

missing property

fed3ca7

remove view() method from array

3f0c1f9

pass device into array view + buffers

3fa3b8d

add reprs

1550e5f

imports

dbabca8

f strings

c4b93d5

docstrings

339d7e5

line length

0fe42cf

fix tests

3395e4b

add comment

f16bb9f

paleolimbot force-pushed the python-device-wrapper branch from 8d61def to f16bb9f Compare November 17, 2023 17:46

jorisvandenbossche reviewed Nov 17, 2023

View reviewed changes

python/src/nanoarrow/_lib.pyx Outdated Show resolved Hide resolved

paleolimbot and others added 3 commits November 17, 2023 14:03

Update python/src/nanoarrow/_lib.pyx

107b8e2

Co-authored-by: Joris Van den Bossche <[email protected]>

explain, fix test

ea6c123

re-fix typo that got lost

065a82e

paleolimbot added 4 commits November 17, 2023 15:03

maybe better test coverage

c44cfb5

better coverage

88e5b1d

fix tests

c4b3cb7

one more cov

966a014

paleolimbot merged commit fda87c5 into apache:main Nov 17, 2023
2 checks passed

paleolimbot deleted the python-device-wrapper branch November 17, 2023 20:33

paleolimbot added this to the nanoarrow 0.4.0 milestone Jan 26, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(python): Add ArrowDeviceArray extension to Python bindings #313

feat(python): Add ArrowDeviceArray extension to Python bindings #313

paleolimbot commented Oct 31, 2023 •

edited

Loading

codecov-commenter commented Oct 31, 2023 •

edited

Loading

paleolimbot commented Nov 15, 2023

jorisvandenbossche left a comment

jorisvandenbossche Nov 16, 2023

paleolimbot Nov 16, 2023

paleolimbot Nov 16, 2023

jorisvandenbossche Nov 16, 2023

jorisvandenbossche Nov 16, 2023

paleolimbot Nov 16, 2023

jorisvandenbossche Nov 16, 2023

jorisvandenbossche Nov 17, 2023

paleolimbot Nov 17, 2023

paleolimbot Nov 17, 2023

feat(python): Add ArrowDeviceArray extension to Python bindings #313

feat(python): Add ArrowDeviceArray extension to Python bindings #313

Conversation

paleolimbot commented Oct 31, 2023 • edited Loading

codecov-commenter commented Oct 31, 2023 • edited Loading

Codecov Report

paleolimbot commented Nov 15, 2023

jorisvandenbossche 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

paleolimbot commented Oct 31, 2023 •

edited

Loading

codecov-commenter commented Oct 31, 2023 •

edited

Loading