pitrou commented on code in PR #42223:
URL: https://github.com/apache/arrow/pull/42223#discussion_r1653314643
##########
python/pyarrow/tests/test_device.py:
##########
@@ -41,3 +41,16 @@ def test_buffer_device():
assert buf.device.is_cpu
assert buf.device == pa.default_cpu_memory_manager().device
assert buf.memory_manager.is_cpu
+
+
+def test_copy_to():
+ mm = pa.default_cpu_memory_manager()
+ arr = pa.array([0, 1, 2])
+ arr_copied = arr.copy_to(mm)
+ assert arr.equals(arr_copied)
+ assert arr.buffers()[1].address != arr_copied.buffers()[1].address
Review Comment:
Can we also check that the device and memory manager are as expected?
##########
python/pyarrow/table.pxi:
##########
@@ -3549,6 +3549,29 @@ cdef class RecordBatch(_Tabular):
row_major, pool))
return pyarrow_wrap_tensor(c_tensor)
+ def copy_to(self, MemoryManager memory_manager):
Review Comment:
We could perhaps also accept a Device and copy to the Device's default
memory manager, for usability.
##########
python/pyarrow/table.pxi:
##########
@@ -3549,6 +3549,29 @@ cdef class RecordBatch(_Tabular):
row_major, pool))
return pyarrow_wrap_tensor(c_tensor)
+ def copy_to(self, MemoryManager memory_manager):
+ """
+ Copy the entire RecordBatch to destination MemoryManager.
+
+ This copies each column of the record batch to create
+ a new record batch where all underlying buffers for the columns have
+ been copied to the destination MemoryManager.
+
+ Parameters
+ ----------
+ memory_manager : pyarrow.MemoryManager
+
+ Returns
+ ------
+ RecordBatch
+ """
+ cdef:
+ shared_ptr[CRecordBatch] c_batch
+
+ with nogil:
+ c_batch =
GetResultValue(self.batch.CopyTo(memory_manager.unwrap()))
Review Comment:
Can `memory_manager` be None here?
--
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]
