feat(python): Clarify interaction between the CDeviceArray, the CArrayView, and the CArray

python/src/nanoarrow/_lib.pyx

@@ -34,7 +34,6 @@ generally have better autocomplete + documentation available to IDEs).
 from libc.stdint cimport uintptr_t, uint8_t, int64_t
 from libc.string cimport memcpy
 from libc.stdio cimport snprintf
-from libc.errno cimport ENOMEM
 from cpython.bytes cimport PyBytes_FromStringAndSize
 from cpython.pycapsule cimport PyCapsule_New, PyCapsule_GetPointer, PyCapsule_IsValid
 from cpython cimport (
@@ -183,7 +182,7 @@ cdef void c_array_shallow_copy(object base, const ArrowArray* c_array,
  c_array_out.release = arrow_array_release
 
 
-cdef object alloc_c_array_shallow_copy(object base, const ArrowArray* c_array) noexcept:
+cdef object alloc_c_array_shallow_copy(object base, const ArrowArray* c_array):
  """Make a shallow copy of an ArrowArray
 
  To more safely implement export of an ArrowArray whose address may be
@@ -198,6 +197,30 @@ cdef object alloc_c_array_shallow_copy(object base, const ArrowArray* c_array) n
  return array_capsule
 
 
+cdef void c_device_array_shallow_copy(object base, const ArrowDeviceArray* c_array,
+ ArrowDeviceArray* c_array_out) noexcept:
+ # shallow copy
+ memcpy(c_array_out, c_array, sizeof(ArrowDeviceArray))
+ c_array_out.array.release = NULL
+ c_array_out.array.private_data = NULL
+
+ # track original base
+ c_array_out.array.private_data = <void*>base
+ Py_INCREF(base)
+ c_array_out.array.release = arrow_array_release
+
+
+cdef object alloc_c_device_array_shallow_copy(object base, const ArrowDeviceArray* c_array):
+ """Make a shallow copy of an ArrowDeviceArray
+
+ See :func:`arrow_c_array_shallow_copy()`
+ """
+ cdef ArrowDeviceArray* c_array_out
+ array_capsule = alloc_c_device_array(&c_array_out)
+ c_device_array_shallow_copy(base, c_array, c_array_out)
+ return array_capsule
+
+
 cdef void pycapsule_buffer_deleter(object stream_capsule) noexcept:
  cdef ArrowBuffer* buffer = <ArrowBuffer*>PyCapsule_GetPointer(
  stream_capsule, 'nanoarrow_buffer'
@@ -207,11 +230,12 @@ cdef void pycapsule_buffer_deleter(object stream_capsule) noexcept:
  ArrowFree(buffer)
 
 
-cdef object alloc_c_buffer(ArrowBuffer** c_buffer) noexcept:
+cdef object alloc_c_buffer(ArrowBuffer** c_buffer):
  c_buffer[0] = <ArrowBuffer*> ArrowMalloc(sizeof(ArrowBuffer))
  ArrowBufferInit(c_buffer[0])
  return PyCapsule_New(c_buffer[0], 'nanoarrow_buffer', &pycapsule_buffer_deleter)
 
+
 cdef void c_deallocate_pybuffer(ArrowBufferAllocator* allocator, uint8_t* ptr, int64_t size) noexcept with gil:
  cdef Py_buffer* buffer = <Py_buffer*>allocator.private_data
  PyBuffer_Release(buffer)
@@ -1027,6 +1051,8 @@ cdef class CArray:
  cdef object _base
  cdef ArrowArray* _ptr
  cdef CSchema _schema
+ cdef ArrowDeviceType _device_type
+ cdef int _device_id
 
  @staticmethod
  def allocate(CSchema schema):
@@ -1038,6 +1064,12 @@ cdef class CArray:
  self._base = base
  self._ptr = <ArrowArray*>addr
  self._schema = schema
+ self._device_type = ARROW_DEVICE_CPU
+ self._device_id = 0
+
+ cdef _set_device(self, ArrowDeviceType device_type, int64_t device_id):
+ self._device_type = device_type
+ self._device_id = device_id
 
  @staticmethod
  def _import_from_c_capsule(schema_capsule, array_capsule):
@@ -1095,7 +1127,9 @@ cdef class CArray:
 
  c_array_out.offset = c_array_out.offset + start
  c_array_out.length = stop - start
- return CArray(base, <uintptr_t>c_array_out, self._schema)
+ cdef CArray out = CArray(base, <uintptr_t>c_array_out, self._schema)
+ out._set_device(self._device_type, self._device_id)
+ return out
 
  def __arrow_c_array__(self, requested_schema=None):
  """
@@ -1115,6 +1149,11 @@ cdef class CArray:
  """
  self._assert_valid()
 
+ if self._device_type != ARROW_DEVICE_CPU:
+ raise ValueError(
+ "Can't invoke __arrow_c_aray__ on non-CPU array "
+ f"with device_type {self._device_type}")
+
  if requested_schema is not None:
  raise NotImplementedError("requested_schema")
 
@@ -1137,10 +1176,22 @@ cdef class CArray:
  if self._ptr.release == NULL:
  raise RuntimeError("CArray is released")
 
+ def view(self):
+ device = CDevice.resolve(self._device_type, self._device_id)
+ return CArrayView.from_array(self, device)
+
  @property
  def schema(self):
  return self._schema
 
+ @property
+ def device_type(self):
+ return self._device_type
+
+ @property
+ def device_id(self):
+ return self._device_id
+
  @property
  def length(self):
  self._assert_valid()
@@ -1175,7 +1226,13 @@ cdef class CArray:
  self._assert_valid()
  if i < 0 or i >= self._ptr.n_children:
  raise IndexError(f"{i} out of range [0, {self._ptr.n_children})")
- return CArray(self._base, <uintptr_t>self._ptr.children[i], self._schema.child(i))
+ cdef CArray out = CArray(
+ self._base,
+ <uintptr_t>self._ptr.children[i],
+ self._schema.child(i)
+ )
+ out._set_device(self._device_type, self._device_id)
+ return out
 
  @property
  def children(self):
@@ -1185,8 +1242,11 @@ cdef class CArray:
  @property
  def dictionary(self):
  self._assert_valid()
+ cdef CArray out
  if self._ptr.dictionary != NULL:
- return CArray(self, <uintptr_t>self._ptr.dictionary, self._schema.dictionary)
+ out = CArray(self, <uintptr_t>self._ptr.dictionary, self._schema.dictionary)
+ out._set_device(self._device_type, self._device_id)
+ return out
  else:
  return None
 
@@ -2330,6 +2390,10 @@ cdef class CDeviceArray:
  self._ptr = <ArrowDeviceArray*>addr
  self._schema = schema
 
+ @property
+ def schema(self):
+ return self._schema
+
  @property
  def device_type(self):
  return self._ptr.device_type
@@ -2340,7 +2404,53 @@ cdef class CDeviceArray:
 
  @property
  def array(self):
- return CArray(self, <uintptr_t>&self._ptr.array, self._schema)
+ # TODO: We loose access to the sync_event here, so we probably need to
+ # synchronize (or propatate it, or somehow prevent data access downstream)
+ cdef CArray array = CArray(self, <uintptr_t>&self._ptr.array, self._schema)
+ array._set_device(self._ptr.device_type, self._ptr.device_id)
+ return array
+
+ def view(self):
+ return self.array.view()
+
+ def __arrow_c_array__(self, requested_schema=None):
+ return self.array.__arrow_c_array__(requested_schema=requested_schema)
+
+ def __arrow_c_device_array__(self, requested_schema=None):
+ if requested_schema is not None:
+ raise NotImplementedError("requested_schema")
+
+ # TODO: evaluate whether we need to synchronize here or whether we should
+ # move device arrays instead of shallow-copying them
+ device_array_capsule = alloc_c_device_array_shallow_copy(self._base, self._ptr)
+ return self._schema.__arrow_c_schema__(), device_array_capsule
+
+ @staticmethod
+ def _import_from_c_capsule(schema_capsule, device_array_capsule):
+ """
+ Import from a ArrowSchema and ArrowArray PyCapsule tuple.
+
+ Parameters
+ ----------
+ schema_capsule : PyCapsule
+ A valid PyCapsule with name 'arrow_schema' containing an
+ ArrowSchema pointer.
+ device_array_capsule : PyCapsule
+ A valid PyCapsule with name 'arrow_device_array' containing an
+ ArrowDeviceArray pointer.
+ """
+ cdef:
+ CSchema out_schema
+ CDeviceArray out
+
+ out_schema = CSchema._import_from_c_capsule(schema_capsule)
+ out = CDeviceArray(
+ device_array_capsule,
+ <uintptr_t>PyCapsule_GetPointer(device_array_capsule, 'arrow_device_array'),
+ out_schema
+ )
+
+ return out
 
  def __repr__(self):
  return _repr_utils.device_array_repr(self)

python/src/nanoarrow/c_lib.py

@@ -427,7 +427,7 @@ def c_array_view(obj, schema=None) -> CArrayView:
  if isinstance(obj, CArrayView) and schema is None:
  return obj
 
- return CArrayView.from_array(c_array(obj, schema))
+ return c_array(obj, schema).view()
 
 
 def c_buffer(obj, schema=None) -> CBuffer:

python/src/nanoarrow/device.py

@@ -16,7 +16,7 @@
 # under the License.
 
 from nanoarrow._lib import CDEVICE_CPU, CDevice, CDeviceArray
-from nanoarrow.c_lib import c_array
+from nanoarrow.c_lib import c_array, c_schema
 
 
 def cpu():
@@ -27,11 +27,20 @@ def resolve(device_type, device_id):
  return CDevice.resolve(device_type, device_id)
 
 
-def c_device_array(obj):
- if isinstance(obj, CDeviceArray):
+def c_device_array(obj, schema=None):
+ if schema is not None:
+ schema = c_schema(schema)
+
+ if isinstance(obj, CDeviceArray) and schema is None:
  return obj
 
- # Only CPU for now
- cpu_array = c_array(obj)
+ if hasattr(obj, "__arrow_c_device_array__"):
+ schema_capsule = None if schema is None else schema.__arrow_c_schema__()
+ schema_capsule, device_array_capsule = obj.__arrow_c_device_array__(
+ requested_schema=schema_capsule
+ )
+ return CDeviceArray._import_from_c_capsule(schema_capsule, device_array_capsule)
 
+ # 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)

python/tests/test_c_array.py

@@ -39,6 +39,8 @@ def test_c_array_from_c_array():
  assert c_array_from_c_array.length == c_array.length
  assert c_array_from_c_array.buffers == c_array.buffers
 
+ assert list(c_array.view().buffer(1)) == [1, 2, 3]
+
 
 def test_c_array_from_capsule_protocol():
  class CArrayWrapper:
@@ -54,6 +56,8 @@ def __arrow_c_array__(self, *args, **kwargs):
  assert c_array_from_protocol.length == c_array.length
  assert c_array_from_protocol.buffers == c_array.buffers
 
+ assert list(c_array_from_protocol.view().buffer(1)) == [1, 2, 3]
+
 
 def test_c_array_from_old_pyarrow():
  # Simulate a pyarrow Array with no __arrow_c_array__
@@ -73,6 +77,8 @@ def _export_to_c(self, *args):
  assert c_array.length == 3
  assert c_array.schema.format == "i"
 
+ assert list(c_array.view().buffer(1)) == [1, 2, 3]
+
  # Make sure that this heuristic won't result in trying to import
  # something else that has an _export_to_c method
  with pytest.raises(TypeError, match="Can't convert object of type DataType"):
@@ -97,6 +103,8 @@ def test_c_array_from_bare_capsule():
  assert c_array_from_capsule.length == c_array.length
  assert c_array_from_capsule.buffers == c_array.buffers
 
+ assert list(c_array_from_capsule.view().buffer(1)) == [1, 2, 3]
+
 
 def test_c_array_type_not_supported():
  with pytest.raises(TypeError, match="Can't convert object of type NoneType"):

python/tests/test_device.py

@@ -17,10 +17,9 @@
 
 import pytest
 
+import nanoarrow as na
 from nanoarrow import device
 
-pa = pytest.importorskip("pyarrow")
-
 
 def test_cpu_device():
  cpu = device.cpu()
@@ -31,12 +30,51 @@ def test_cpu_device():
  cpu = device.resolve(1, 0)
  assert cpu.device_type == 1
 
- pa_array = pa.array([1, 2, 3])
 
- darray = device.c_device_array(pa_array)
+def test_c_device_array():
+ # Unrecognized arguments should be passed to c_array() to generate CPU array
+ darray = device.c_device_array([1, 2, 3], na.int32())
+
  assert darray.device_type == 1
  assert darray.device_id == 0
- assert darray.array.length == 3
  assert "device_type: 1" in repr(darray)
 
+ assert darray.schema.format == "i"
+
+ assert darray.array.length == 3
+ assert darray.array.device_type == device.cpu().device_type
+ assert darray.array.device_id == device.cpu().device_id
+
+ darray_view = darray.view()
+ assert darray_view.length == 3
+ assert list(darray_view.buffer(1)) == [1, 2, 3]
+
+ # A CDeviceArray should be returned as is
  assert device.c_device_array(darray) is darray
+
+ # A CPU device array should be able to export to a regular array
+ array = na.c_array(darray)
+ assert array.schema.format == "i"
+ assert array.buffers == darray.array.buffers
+
+
+def test_c_device_array_protocol():
+ # Wrapper to prevent c_device_array() from returning early when it detects the
+ # input is already a CDeviceArray
+ class CDeviceArrayWrapper:
+ 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)
+
+ darray = device.c_device_array([1, 2, 3], na.int32())
+ wrapper = CDeviceArrayWrapper(darray)
+
+ darray2 = device.c_device_array(wrapper)
+ assert darray2.schema.format == "i"
+ assert darray2.array.length == 3
+ assert darray2.array.buffers == darray.array.buffers
+
+ with pytest.raises(NotImplementedError):
+ device.c_device_array(wrapper, na.int64())