From a84219e15ce20cdc6ff59667639d9515fe32ec9d Mon Sep 17 00:00:00 2001 From: anjakefala Date: Thu, 25 Jul 2024 21:44:37 +0000 Subject: [PATCH] WIP Buffer Copy binding --- python/pyarrow/device.pxi | 4 +++- python/pyarrow/includes/libarrow.pxd | 1 + python/pyarrow/io.pxi | 12 ++++++++++++ python/pyarrow/lib.pxd | 2 ++ 4 files changed, 18 insertions(+), 1 deletion(-) diff --git a/python/pyarrow/device.pxi b/python/pyarrow/device.pxi index 6e6034752085a..6ca5761893ed1 100644 --- a/python/pyarrow/device.pxi +++ b/python/pyarrow/device.pxi @@ -130,6 +130,9 @@ cdef class MemoryManager(_Weakrefable): self.init(mm) return self + cdef inline shared_ptr[CMemoryManager] unwrap(self) nogil: + return self.memory_manager + def __repr__(self): return "".format( frombytes(self.memory_manager.get().device().get().ToString()) @@ -152,7 +155,6 @@ cdef class MemoryManager(_Weakrefable): """ return self.memory_manager.get().is_cpu() - def default_cpu_memory_manager(): """ Return the default CPU MemoryManager instance. diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 0d871f411b11b..277e5f8dacd3c 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -364,6 +364,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: shared_ptr[CDevice] device() const shared_ptr[CMemoryManager] memory_manager() CDeviceAllocationType device_type() + CResult[shared_ptr[CBuffer]] Copy(const shared_ptr[CMemoryManager]& mm) const CResult[shared_ptr[CBuffer]] SliceBufferSafe( const shared_ptr[CBuffer]& buffer, int64_t offset) diff --git a/python/pyarrow/io.pxi b/python/pyarrow/io.pxi index 1d942e8ccabc6..71f9a5714d0a9 100644 --- a/python/pyarrow/io.pxi +++ b/python/pyarrow/io.pxi @@ -1446,6 +1446,18 @@ cdef class Buffer(_Weakrefable): """ return _wrap_device_allocation_type(self.buffer.get().device_type()) + def copy(self, MemoryManager mm): + """ + The buffer contents will be copied into a new buffer allocated by the + given MemoryManager. This function supports cross-device copies. + """ + cdef: + shared_ptr[CBuffer] c_buffer + + c_buffer = GetResultValue(self.buffer.get().Copy(mm.unwrap())) + 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/lib.pxd b/python/pyarrow/lib.pxd index 082d8470cdbb0..bedcff9559d4e 100644 --- a/python/pyarrow/lib.pxd +++ b/python/pyarrow/lib.pxd @@ -544,6 +544,8 @@ cdef class MemoryManager(_Weakrefable): @staticmethod cdef wrap(const shared_ptr[CMemoryManager]& mm) + cdef inline shared_ptr[CMemoryManager] unwrap(self) nogil + cdef class Buffer(_Weakrefable): cdef: