diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx index e7d5ed7b8bfec..fdfc37cc6aa84 100644 --- a/python/pyarrow/_cuda.pyx +++ b/python/pyarrow/_cuda.pyx @@ -185,10 +185,27 @@ cdef class Context(_Weakrefable): cudabuf = GetResultValue(self.context.get().Allocate(nbytes)) return pyarrow_wrap_cudabuffer(cudabuf) + @property def memory_manager(self): + """ + The default memory manager tied to this context's device. + + Returns + ------- + MemoryManager + """ return MemoryManager.wrap(self.context.get().memory_manager()) + @property def device(self): + """ + The device instance associated with this context. + + + Returns + ------- + Device + """ return Device.wrap(self.context.get().device()) def foreign_buffer(self, address, size, base=None): diff --git a/python/pyarrow/includes/libarrow_cuda.pxd b/python/pyarrow/includes/libarrow_cuda.pxd index 39ca53908a68b..7019d3483dc74 100644 --- a/python/pyarrow/includes/libarrow_cuda.pxd +++ b/python/pyarrow/includes/libarrow_cuda.pxd @@ -41,8 +41,8 @@ cdef extern from "arrow/gpu/cuda_api.h" namespace "arrow::cuda" nogil: const void* handle() const int device_number() const CResult[uintptr_t] GetDeviceAddress(uintptr_t addr) - shared_ptr[CDevice] device() const - shared_ptr[CMemoryManager] memory_manager() const + shared_ptr[CDevice] device const + shared_ptr[CMemoryManager] memory_manager const cdef cppclass CCudaIpcMemHandle" arrow::cuda::CudaIpcMemHandle": @staticmethod diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 021a5bfb9d614..2644cfe42c198 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -67,7 +67,6 @@ def test_Context(): assert not dev.is_cpu assert dev.device_id == 0 assert dev.device_type == pa.DeviceAllocationType.CUDA - assert "