From 17edaa934fbea80df18f5cc48f17af46bac6c6be Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Thu, 20 Jun 2024 14:30:22 +0000 Subject: [PATCH] GH-41126: [Python] Test Buffer device/device_type access on CUDA --- python/pyarrow/tests/test_cuda.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 400db2643bd56..1a287ecaff486 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -534,6 +534,25 @@ def put(*args, **kwargs): put(position=position, nbytes=nbytes) +def test_buffer_device(): + buf = cuda.new_host_buffer(10) + assert buf.device_type == pa.DeviceAllocationType.CUDA_HOST + assert isinstance(buf.device, pa.Device) + assert isinstance(buf.memory_manager, pa.MemoryManager) + assert buf.is_cpu + assert buf.device.is_cpu + assert buf.device == pa.default_cpu_memory_manager().device + assert buf.memory_manager.is_cpu + + _, buf = make_random_buffer(size=10, target='device') + assert buf.device_type == pa.DeviceAllocationType.CUDA + assert isinstance(buf.device, pa.Device) + assert isinstance(buf.memory_manager, pa.MemoryManager) + assert not buf.is_cpu + assert not buf.device.is_cpu + assert not buf.memory_manager.is_cpu + + def test_BufferWriter(): def allocate(size): cbuf = global_context.new_buffer(size)