diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 8e6922a912a..8bd29101ec6 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -367,6 +367,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: const shared_ptr[CMemoryManager] memory_manager() CDeviceAllocationType device_type() + @staticmethod + CResult[shared_ptr[CBuffer]] Copy(shared_ptr[CBuffer] source, const shared_ptr[CMemoryManager]& to) + CResult[shared_ptr[CBuffer]] SliceBufferSafe( const shared_ptr[CBuffer]& buffer, int64_t offset) CResult[shared_ptr[CBuffer]] SliceBufferSafe( diff --git a/python/pyarrow/io.pxi b/python/pyarrow/io.pxi index 1d942e8ccab..65c4485e55f 100644 --- a/python/pyarrow/io.pxi +++ b/python/pyarrow/io.pxi @@ -1446,6 +1446,41 @@ 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 = (destination).unwrap().get().default_memory_manager() + elif isinstance(destination, MemoryManager): + c_memory_manager = (destination).unwrap() + else: + raise TypeError( + "Argument 'destination' is incorrect type (expected pyarrow.Device or " + f"pyarrow.MemoryManager, got {type(destination)})" + ) + + with nogil: + c_buffer = GetResultValue(CBuffer.Copy(self.buffer, c_memory_manager)) + return pyarrow_wrap_buffer(c_buffer) + @property def parent(self): cdef shared_ptr[CBuffer] parent_buf = self.buffer.get().parent() diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index a71fa036503..ec28e4ca731 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -119,6 +119,25 @@ 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 + assert buf2.copy(pa.default_cpu_memory_manager()).equals(buf) + cudabuf2 = cuda.CudaBuffer.from_buffer(buf2) + assert cudabuf2.size == cudabuf.size + assert not cudabuf2.is_cpu + assert cudabuf2.device_type == pa.DeviceAllocationType.CUDA + + assert cudabuf2.copy_to_host().equals(buf) + + assert cudabuf2.device == mm2.device + + @pytest.mark.parametrize("size", [0, 1, 1000]) def test_context_device_buffer(size): # Creating device buffer from host buffer; diff --git a/python/pyarrow/tests/test_io.py b/python/pyarrow/tests/test_io.py index e2df1b1c468..a0fc697fdf0 100644 --- a/python/pyarrow/tests/test_io.py +++ b/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 = buf_on_gpu.copy(dest) + cuda_buf2 = cuda.CudaBuffer.from_buffer(buf2) + assert cuda_buf2.size == cuda_buf.size + assert cuda_buf2.to_pybytes() == b'testing' + assert buf2.device == mm.device + # Buffers on different devices msg_device = "Device on which the data resides differs between buffers" with pytest.raises(ValueError, match=msg_device):