jorisvandenbossche commented on code in PR #550:
URL: https://github.com/apache/arrow-nanoarrow/pull/550#discussion_r1691207145
##########
python/src/nanoarrow/_array.pyx:
##########
@@ -107,7 +106,7 @@ cdef class CArrayView:
def __cinit__(self, object base, uintptr_t addr):
self._base = base
self._ptr = <ArrowArrayView*>addr
- self._device = DEVICE_CPU
+ self._event = CSharedSyncEvent(DEVICE_CPU)
Review Comment:
Should we make `event` an option param in the init here? (which still
defaults to the above if not specified)
(to avoid having to assign like `child._event = self._event` after
constructing the new CArrayView, like is done below in the `dictionary`
property)
EDIT: although I see now that for CArray this is also done after the init,
but with a `_set_device` helper
##########
python/src/nanoarrow/_array.pyx:
##########
@@ -747,6 +766,16 @@ cdef class CArrayBuilder:
return out
+ def finish_device(self):
Review Comment:
Can you add a docstring here?
##########
python/src/nanoarrow/_buffer.pyx:
##########
@@ -163,15 +193,31 @@ cdef DLDevice view_to_dlpack_device(CBufferView view):
raise ValueError('DataType is not compatible with DLPack spec: ' +
view.data_type)
# Define DLDevice struct
- if view._device.device_type_id == ARROW_DEVICE_CPU:
+ cdef ArrowDevice* arrow_device = view._event.device._ptr
+ if arrow_device.device_type is ARROW_DEVICE_CPU:
+ # DLPack uses 0 for the CPU device id where Arrow uses -1
device.device_type = kDLCPU
device.device_id = 0
else:
- raise ValueError('Only CPU device is currently supported.')
+ # Otherwise, Arrow's device identifiers and types are intentionally
+ # identical to DLPack
+ device.device_type = <DLDeviceType>arrow_device.device_type
+ device.device_id = arrow_device.device_id
Review Comment:
This enables exporting device data through dlpack, right? But should we then
still handle the sync event somehow elsewhere?
##########
python/src/nanoarrow/_array.pyx:
##########
@@ -107,7 +106,7 @@ cdef class CArrayView:
def __cinit__(self, object base, uintptr_t addr):
self._base = base
self._ptr = <ArrowArrayView*>addr
- self._device = DEVICE_CPU
+ self._event = CSharedSyncEvent(DEVICE_CPU)
Review Comment:
And you are replacing device with event here because that's the only thing
that in practice is needed in methods from the ArrayView class?
##########
python/src/nanoarrow/_array.pyx:
##########
@@ -669,6 +685,9 @@ cdef class CArrayBuilder:
if i < 0 or i > 3:
raise IndexError("i must be >= 0 and <= 3")
+ if buffer._device != self._device:
+ raise ValueError("Can't source and buffer device not identical")
Review Comment:
Some grammatical error in this sentence I think (just remove the "can't"?)
##########
python/src/nanoarrow/_array.pyx:
##########
@@ -636,6 +648,10 @@ cdef class CArrayBuilder:
self._ptr.null_count = 0
return self
+ # Don't attempt to access a non-cpu buffer
+ if self._device != DEVICE_CPU:
Review Comment:
Maybe update the docstring above that this is only fully synchronized for
CPU buffers?
##########
python/src/nanoarrow/c_array.py:
##########
@@ -298,7 +303,10 @@ def c_array_from_buffers(
builder.resolve_null_count()
# Validate + finish
- return builder.finish(validation_level=validation_level)
+ if device is DEVICE_CPU:
+ return builder.finish(validation_level=validation_level)
+ else:
+ return builder.finish_device()
Review Comment:
This means it returns a ArrowArray vs ArrowDeviceArray depending on the
device? (but you might also want a device array for cpu device?)
##########
python/src/nanoarrow/_buffer.pyx:
##########
@@ -569,6 +618,36 @@ cdef class CBuffer:
out._populate_view()
return out
+ @staticmethod
+ def from_dlpack(obj):
+ capsule = obj.__dlpack__()
Review Comment:
Do we need to pass `stream` here in case of CUDA device?
##########
python/src/nanoarrow/c_array.py:
##########
@@ -201,6 +202,7 @@ def c_array_from_buffers(
children: Iterable[Any] = (),
validation_level: Literal[None, "full", "default", "minimal", "none"] =
None,
move: bool = False,
+ device: Union[Device, None] = None,
Review Comment:
Can you add this keyword to the docstring?
##########
python/src/nanoarrow/_utils.pyx:
##########
@@ -493,3 +493,29 @@ cdef object c_buffer_set_pybuffer(object obj,
ArrowBuffer** c_buffer):
# Return the calculated components
return format
+
+
+cdef void c_deallocate_pyobject(ArrowBufferAllocator* allocator, uint8_t* ptr,
int64_t size) noexcept with gil:
+ """ArrowBufferDeallocatorCallback for an ArrowBuffer wrapping a
Py_Buffer"""
+ Py_DECREF(<object>allocator.private_data)
+
+
+cdef ArrowBufferAllocator c_pyobject_deallocator(object obj):
+ """ArrowBufferAllocator implementation wrapping a PyObject"""
+ Py_INCREF(obj)
+ return ArrowBufferDeallocator(
+ <ArrowBufferDeallocatorCallback>&c_deallocate_pyobject,
+ <void*>obj
+ )
+
+cdef void c_buffer_set_pyobject(object base, uint8_t* data, int64_t
size_bytes, ArrowBuffer** c_buffer):
+ """Manage a Py_Buffer reference as an ArrowBuffer
Review Comment:
This is copy pasted from above and not fully correct here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
