Skip to content

Commit

Permalink
Apply PR comments
Browse files Browse the repository at this point in the history
  • Loading branch information
anjakefala committed Sep 13, 2024
1 parent 5ad2d31 commit 2e10905
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 10 deletions.
4 changes: 3 additions & 1 deletion python/pyarrow/includes/libarrow.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -366,7 +366,9 @@ 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(shared_ptr[CBuffer] source, const shared_ptr[CMemoryManager]& to) const

@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)
Expand Down
20 changes: 11 additions & 9 deletions python/pyarrow/io.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -1446,15 +1446,17 @@ cdef class Buffer(_Weakrefable):
"""
return _wrap_device_allocation_type(self.buffer.get().device_type())

def copy(self, dest):
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
----------
dest : pyarrow.MemoryManager or pyarrow.Device
destination : pyarrow.MemoryManager or pyarrow.Device
Used to allocate the new buffer.
Returns
Expand All @@ -1465,17 +1467,17 @@ cdef class Buffer(_Weakrefable):
shared_ptr[CBuffer] c_buffer
shared_ptr[CMemoryManager] c_memory_manager

if isinstance(dest, Device):
c_memory_manager = (<Device>dest).unwrap().get().default_memory_manager()
elif isinstance(dest, MemoryManager):
c_memory_manager = (<MemoryManager>dest).unwrap()
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 'dest' is incorrect type (expected pyarrow.Device or "
f"pyarrow.MemoryManager, got {type(dest)})"
"Argument 'destination' is incorrect type (expected pyarrow.Device or "
f"pyarrow.MemoryManager, got {type(destination)})"
)

c_buffer = GetResultValue(self.buffer.get().Copy(self.buffer, c_memory_manager))
c_buffer = GetResultValue(CBuffer.Copy(self.buffer, c_memory_manager))
return pyarrow_wrap_buffer(c_buffer)

@property
Expand Down

0 comments on commit 2e10905

Please sign in to comment.