diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 8bd29101ec621..3c9b43b50b03d 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -370,6 +370,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: @staticmethod CResult[shared_ptr[CBuffer]] Copy(shared_ptr[CBuffer] source, const shared_ptr[CMemoryManager]& to) + CResult[shared_ptr[CBuffer]] View(shared_ptr[CBuffer] source, const shared_ptr[CMemoryManager]& to) const + 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 646b33b142770..5564f7845d729 100644 --- a/python/pyarrow/io.pxi +++ b/python/pyarrow/io.pxi @@ -1480,6 +1480,47 @@ cdef class Buffer(_Weakrefable): c_buffer = GetResultValue(CBuffer.Copy(self.buffer, c_memory_manager)) return pyarrow_wrap_buffer(c_buffer) + def view(self, destination): + """ + Zero-copy access this buffer, from destination. + + The underlying mechanism is typically implemented by the kernel or device driver, + and may involve lazy caching of parts of the buffer contents on the destination + device's memory. + + If a non-copy view is unsupported for the buffer on the given device, + an error is raised. + + Parameters + ---------- + destination : pyarrow.MemoryManager or pyarrow.Device + MemoryManager or device on which to view buffer contents + + Returns + ------- + Buffer or None + """ + 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)})" + ) + + c_buffer = GetResultValue(self.buffer.get().View(self.buffer, c_memory_manager)) + + if c_buffer.get() == nullptr: + return None + + 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 9e9d543ee1cea..b86d9f595fc69 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -137,6 +137,28 @@ def test_copy_from_buffer(): assert cudabuf2.device == mm2.device +def test_view_of_buffer(): + data = np.array([1, 2, 3, 4, 5]) + buf = pa.py_buffer(data) + cudabuf = global_context.buffer_from_data(buf) + + mm2 = global_context1.memory_manager + for dest in [mm2, mm2.device]: + buf2 = cudabuf.view(dest) + assert buf2.device_type == pa.DeviceAllocationType.CUDA + 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 + + # test that it is a view, and not a copy + data[1] = 20 + assert np.frombuffer(buf2, dtype=np.int32).tolist() == [20, 2, 3, 4, 5] + + @pytest.mark.parametrize("size", [0, 1, 1000]) def test_context_device_buffer(size): # Creating device buffer from host buffer;