pitrou commented on code in PR #43590:
URL: https://github.com/apache/arrow/pull/43590#discussion_r1764717499
##########
python/pyarrow/tests/test_io.py:
##########
@@ -710,6 +710,15 @@ def test_non_cpu_buffer(pickle_module):
# Sliced buffers with same address
assert buf_on_gpu_sliced.equals(cuda_buf[2:4])
+ # Testing copying buffer
+ mm = ctx.memory_manager
+ for dest in [mm, mm.device]:
+ buf2 = cuda_buf.copy(dest)
+ cuda_buf2 = cuda.CudaBuffer.from_buffer(buf2)
+ cuda_buf2.size == cuda_buf.size
+ cuda_buf2.copy_to_host()[:] == b'testing'
Review Comment:
Are you missing `assert`s here?
##########
python/pyarrow/io.pxi:
##########
@@ -1446,6 +1446,40 @@ cdef class Buffer(_Weakrefable):
"""
return _wrap_device_allocation_type(self.buffer.get().device_type())
+ def copy(self, destination):
+ """
+ Copy the buffer to the destination device.
+
+ The buffer contents will be copied into a new buffer allocated onto
+ the destination MemoryManager or Device. It will return the new Buffer.
+ This function supports cross-device copies.
+
+ Parameters
+ ----------
+ destination : pyarrow.MemoryManager or pyarrow.Device
+ Used to allocate the new buffer.
+
+ Returns
+ -------
+ Buffer
+ """
+ cdef:
+ shared_ptr[CBuffer] c_buffer
+ shared_ptr[CMemoryManager] c_memory_manager
+
+ if isinstance(destination, Device):
+ c_memory_manager =
(<Device>destination).unwrap().get().default_memory_manager()
+ elif isinstance(destination, MemoryManager):
+ c_memory_manager = (<MemoryManager>destination).unwrap()
+ else:
+ raise TypeError(
+ "Argument 'destination' is incorrect type (expected
pyarrow.Device or "
+ f"pyarrow.MemoryManager, got {type(destination)})"
+ )
+
+ c_buffer = GetResultValue(CBuffer.Copy(self.buffer, c_memory_manager))
Review Comment:
This should release the GIL since copying is an expensive operation.
##########
python/pyarrow/tests/test_cuda.py:
##########
@@ -119,6 +119,24 @@ def make_random_buffer(size, target='host'):
raise ValueError('invalid target value')
+def test_copy_from_buffer():
+ arr, buf = make_random_buffer(128)
+ cudabuf = global_context.buffer_from_data(buf)
+
+ mm2 = global_context1.memory_manager
+ for dest in [mm2, mm2.device]:
+ buf2 = cudabuf.copy(dest)
+ assert buf2.device_type == pa.DeviceAllocationType.CUDA
Review Comment:
Perhaps also try copying directly to the CPU?
```suggestion
assert buf2.device_type == pa.DeviceAllocationType.CUDA
assert buf2.copy(pa.default_memory_manager()).equals(buf)
```
--
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]
