From 931954349a0a558d983958915a7aaa87cb00186d Mon Sep 17 00:00:00 2001 From: anjakefala Date: Thu, 19 Sep 2024 19:03:52 +0000 Subject: [PATCH] Antoine review comments --- python/pyarrow/io.pxi | 3 ++- python/pyarrow/tests/test_cuda.py | 1 + python/pyarrow/tests/test_io.py | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/python/pyarrow/io.pxi b/python/pyarrow/io.pxi index 646b33b142770..65c4485e55fff 100644 --- a/python/pyarrow/io.pxi +++ b/python/pyarrow/io.pxi @@ -1477,7 +1477,8 @@ cdef class Buffer(_Weakrefable): f"pyarrow.MemoryManager, got {type(destination)})" ) - c_buffer = GetResultValue(CBuffer.Copy(self.buffer, c_memory_manager)) + with nogil: + c_buffer = GetResultValue(CBuffer.Copy(self.buffer, c_memory_manager)) return pyarrow_wrap_buffer(c_buffer) @property diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 9e9d543ee1cea..ec28e4ca73137 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -127,6 +127,7 @@ def test_copy_from_buffer(): 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 diff --git a/python/pyarrow/tests/test_io.py b/python/pyarrow/tests/test_io.py index 19d68bb674a3b..7889badd074d3 100644 --- a/python/pyarrow/tests/test_io.py +++ b/python/pyarrow/tests/test_io.py @@ -715,8 +715,8 @@ def test_non_cpu_buffer(pickle_module): for dest in [mm, mm.device]: buf2 = cuda_buf.copy(dest) cuda_buf2 = cuda.CudaBuffer.from_buffer(buf2) - cuda_buf2.size == cuda_buf.size - cuda_buf2.copy_to_host()[:] == b'testing' + assert cuda_buf2.size == cuda_buf.size + assert cuda_buf2.copy_to_host()[:] == b'testing' assert cuda_buf2.device == mm.device # Buffers on different devices