refactor(python): Split device functionality into its own module (#548)

Similar to previous PRs, this PR moves the device class definition into
its own module. Both arrays and buffers are device-aware, so splitting
the device definition is a prerequisite.

Author: paleolimbot

python/setup.py

@@ -95,6 +95,21 @@ def get_version(pkg_path):
 
 setup(
     ext_modules=[
+        Extension(
+            name="nanoarrow._device",
+            include_dirs=["src/nanoarrow", "vendor"],
+            language="c",
+            sources=[
+                "src/nanoarrow/_device.pyx",
+                "vendor/nanoarrow.c",
+                "vendor/nanoarrow_device.c",
+            ],
+            extra_compile_args=extra_compile_args,
+            extra_link_args=extra_link_args,
+            define_macros=extra_define_macros,
+            library_dirs=library_dirs,
+            libraries=libraries,
+        ),
         Extension(
             name="nanoarrow._types",
             include_dirs=["src/nanoarrow", "vendor"],

python/src/nanoarrow/_device.pxd

@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# cython: language_level = 3
+
+from nanoarrow_device_c cimport ArrowDevice
+
+cdef class Device:
+    cdef object _base
+    cdef ArrowDevice* _ptr

python/src/nanoarrow/_device.pyx

@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# cython: language_level = 3
+
+from libc.stdint cimport uintptr_t, int64_t
+
+from nanoarrow_device_c cimport (
+    ARROW_DEVICE_CPU,
+    ARROW_DEVICE_CUDA,
+    ARROW_DEVICE_CUDA_HOST,
+    ARROW_DEVICE_OPENCL,
+    ARROW_DEVICE_VULKAN,
+    ARROW_DEVICE_METAL,
+    ARROW_DEVICE_VPI,
+    ARROW_DEVICE_ROCM,
+    ARROW_DEVICE_ROCM_HOST,
+    ARROW_DEVICE_EXT_DEV,
+    ARROW_DEVICE_CUDA_MANAGED,
+    ARROW_DEVICE_ONEAPI,
+    ARROW_DEVICE_WEBGPU,
+    ARROW_DEVICE_HEXAGON,
+    ArrowDevice,
+    ArrowDeviceCpu,
+    ArrowDeviceResolve
+)
+
+from enum import Enum
+
+from nanoarrow import _repr_utils
+
+
+class DeviceType(Enum):
+    """
+    An enumerator providing access to the device constant values
+    defined in the Arrow C Device interface. Unlike the other enum
+    accessors, this Python Enum is defined in Cython so that we can use
+    the bulit-in functionality to do better printing of device identifiers
+    for classes defined in Cython. Unlike the other enums, users don't
+    typically need to specify these (but would probably like them printed
+    nicely).
+    """
+
+    CPU = ARROW_DEVICE_CPU
+    CUDA = ARROW_DEVICE_CUDA
+    CUDA_HOST = ARROW_DEVICE_CUDA_HOST
+    OPENCL = ARROW_DEVICE_OPENCL
+    VULKAN =  ARROW_DEVICE_VULKAN
+    METAL = ARROW_DEVICE_METAL
+    VPI = ARROW_DEVICE_VPI
+    ROCM = ARROW_DEVICE_ROCM
+    ROCM_HOST = ARROW_DEVICE_ROCM_HOST
+    EXT_DEV = ARROW_DEVICE_EXT_DEV
+    CUDA_MANAGED = ARROW_DEVICE_CUDA_MANAGED
+    ONEAPI = ARROW_DEVICE_ONEAPI
+    WEBGPU = ARROW_DEVICE_WEBGPU
+    HEXAGON = ARROW_DEVICE_HEXAGON
+
+
+cdef class Device:
+    """ArrowDevice wrapper
+
+    The ArrowDevice structure is a nanoarrow internal struct (i.e.,
+    not ABI stable) that contains callbacks for device operations
+    beyond its type and identifier (e.g., copy buffers to or from
+    a device).
+    """
+
+    def __cinit__(self, object base, uintptr_t addr):
+        self._base = base,
+        self._ptr = <ArrowDevice*>addr
+
+    def __repr__(self):
+        return _repr_utils.device_repr(self)
+
+    @property
+    def device_type(self):
+        return DeviceType(self._ptr.device_type)
+
+    @property
+    def device_type_id(self):
+        return self._ptr.device_type
+
+    @property
+    def device_id(self):
+        return self._ptr.device_id
+
+    @staticmethod
+    def resolve(device_type, int64_t device_id):
+        if int(device_type) == ARROW_DEVICE_CPU:
+            return DEVICE_CPU
+
+        cdef ArrowDevice* c_device = ArrowDeviceResolve(device_type, device_id)
+        if c_device == NULL:
+            raise ValueError(f"Device not found for type {device_type}/{device_id}")
+
+        return Device(None, <uintptr_t>c_device)
+
+
+# Cache the CPU device
+# The CPU device is statically allocated (so base is None)
+DEVICE_CPU = Device(None, <uintptr_t>ArrowDeviceCpu())

python/src/nanoarrow/_lib.pyx

@@ -53,9 +53,7 @@ from nanoarrow_c cimport *
 from nanoarrow_device_c cimport *
 from nanoarrow_dlpack cimport *
 
-from enum import Enum
-from struct import unpack_from, iter_unpack, calcsize, Struct
-from nanoarrow import _repr_utils
+from nanoarrow._device cimport Device
 
 from nanoarrow cimport _types
 from nanoarrow._utils cimport (
@@ -71,6 +69,11 @@ from nanoarrow._utils cimport (
     Error
 )
 
+from struct import unpack_from, iter_unpack, calcsize, Struct
+from nanoarrow import _repr_utils
+from nanoarrow._device import DEVICE_CPU, DeviceType
+
+
 cdef void pycapsule_dlpack_deleter(object dltensor) noexcept:
     cdef DLManagedTensor* dlm_tensor
 
@@ -176,91 +179,6 @@ cdef class CArrowTimeUnit:
     NANO = NANOARROW_TIME_UNIT_NANO
 
 
-class DeviceType(Enum):
-    """
-    An enumerator providing access to the device constant values
-    defined in the Arrow C Device interface. Unlike the other enum
-    accessors, this Python Enum is defined in Cython so that we can use
-    the bulit-in functionality to do better printing of device identifiers
-    for classes defined in Cython. Unlike the other enums, users don't
-    typically need to specify these (but would probably like them printed
-    nicely).
-    """
-
-    CPU = ARROW_DEVICE_CPU
-    CUDA = ARROW_DEVICE_CUDA
-    CUDA_HOST = ARROW_DEVICE_CUDA_HOST
-    OPENCL = ARROW_DEVICE_OPENCL
-    VULKAN =  ARROW_DEVICE_VULKAN
-    METAL = ARROW_DEVICE_METAL
-    VPI = ARROW_DEVICE_VPI
-    ROCM = ARROW_DEVICE_ROCM
-    ROCM_HOST = ARROW_DEVICE_ROCM_HOST
-    EXT_DEV = ARROW_DEVICE_EXT_DEV
-    CUDA_MANAGED = ARROW_DEVICE_CUDA_MANAGED
-    ONEAPI = ARROW_DEVICE_ONEAPI
-    WEBGPU = ARROW_DEVICE_WEBGPU
-    HEXAGON = ARROW_DEVICE_HEXAGON
-
-
-cdef class Device:
-    """ArrowDevice wrapper
-
-    The ArrowDevice structure is a nanoarrow internal struct (i.e.,
-    not ABI stable) that contains callbacks for device operations
-    beyond its type and identifier (e.g., copy buffers to or from
-    a device).
-    """
-
-    cdef object _base
-    cdef ArrowDevice* _ptr
-
-    def __cinit__(self, object base, uintptr_t addr):
-        self._base = base,
-        self._ptr = <ArrowDevice*>addr
-
-    def _array_init(self, uintptr_t array_addr, CSchema schema):
-        cdef ArrowArray* array_ptr = <ArrowArray*>array_addr
-        cdef ArrowDeviceArray* device_array_ptr
-        cdef void* sync_event = NULL
-        holder = alloc_c_device_array(&device_array_ptr)
-        cdef int code = ArrowDeviceArrayInit(self._ptr, device_array_ptr, array_ptr, sync_event)
-        Error.raise_error_not_ok("ArrowDevice::init_array", code)
-
-        return CDeviceArray(holder, <uintptr_t>device_array_ptr, schema)
-
-    def __repr__(self):
-        return _repr_utils.device_repr(self)
-
-    @property
-    def device_type(self):
-        return DeviceType(self._ptr.device_type)
-
-    @property
-    def device_type_id(self):
-        return self._ptr.device_type
-
-    @property
-    def device_id(self):
-        return self._ptr.device_id
-
-    @staticmethod
-    def resolve(device_type, int64_t device_id):
-        if int(device_type) == ARROW_DEVICE_CPU:
-            return DEVICE_CPU
-
-        cdef ArrowDevice* c_device = ArrowDeviceResolve(device_type, device_id)
-        if c_device == NULL:
-            raise ValueError(f"Device not found for type {device_type}/{device_id}")
-
-        return Device(None, <uintptr_t>c_device)
-
-
-# Cache the CPU device
-# The CPU device is statically allocated (so base is None)
-DEVICE_CPU = Device(None, <uintptr_t>ArrowDeviceCpu())
-
-
 cdef class CSchema:
     """Low-level ArrowSchema wrapper
 
@@ -2808,6 +2726,17 @@ cdef class CDeviceArray:
         self._ptr = <ArrowDeviceArray*>addr
         self._schema = schema
 
+    @staticmethod
+    def _init_from_array(Device device, uintptr_t array_addr, CSchema schema):
+        cdef ArrowArray* array_ptr = <ArrowArray*>array_addr
+        cdef ArrowDeviceArray* device_array_ptr
+        cdef void* sync_event = NULL
+        holder = alloc_c_device_array(&device_array_ptr)
+        cdef int code = ArrowDeviceArrayInit(device._ptr, device_array_ptr, array_ptr, sync_event)
+        Error.raise_error_not_ok("ArrowDeviceArrayInit", code)
+
+        return CDeviceArray(holder, <uintptr_t>device_array_ptr, schema)
+
     @property
     def schema(self):
         return self._schema

python/src/nanoarrow/array.py

@@ -19,14 +19,8 @@
 from functools import cached_property
 from typing import Iterable, Tuple
 
-from nanoarrow._lib import (
-    DEVICE_CPU,
-    CArray,
-    CArrayView,
-    CBuffer,
-    CMaterializedArrayStream,
-    Device,
-)
+from nanoarrow._device import DEVICE_CPU, Device
+from nanoarrow._lib import CArray, CArrayView, CBuffer, CMaterializedArrayStream
 from nanoarrow.c_array import c_array, c_array_view
 from nanoarrow.c_array_stream import c_array_stream
 from nanoarrow.c_schema import c_schema

python/src/nanoarrow/device.py

@@ -15,7 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from nanoarrow._lib import DEVICE_CPU, CDeviceArray, Device, DeviceType  # noqa: F401
+from nanoarrow._device import DEVICE_CPU, Device, DeviceType  # noqa: F401
+from nanoarrow._lib import CDeviceArray
 from nanoarrow.c_array import c_array
 from nanoarrow.c_schema import c_schema
 
@@ -44,4 +45,4 @@ def c_device_array(obj, schema=None):
 
     # Attempt to create a CPU array and wrap it
     cpu_array = c_array(obj, schema=schema)
-    return cpu()._array_init(cpu_array._addr(), cpu_array.schema)
+    return CDeviceArray._init_from_array(cpu(), cpu_array._addr(), cpu_array.schema)