diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 8749ab29d08..36b97a62064 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -545,6 +545,28 @@ 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 + # it is not entirely clear if CudaHostBuffer should use the default CPU memory + # manager (as it does now), see https://github.com/apache/arrow/pull/42221 + 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 buf.device == global_context.memory_manager.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)