GH-41126: [Python] Basic bindings for Device and MemoryManager classes #41685
Conversation
|
|
github-actions
bot
added
Component: Python
awaiting committer review
| """ | ||
| Return the DeviceAllocationType of this device. | ||
| """ | ||
| return _wrap_device_allocation_type(self.device.get().device_type()) |
There was a problem hiding this comment.
Why do we wrap device_type, but not other properties like is_cpu?
There was a problem hiding this comment.
is_cpu is a simple bool, and cython takes care of returning that as a python object. For the type here I wanted to return a python enum (and not just an integer)
There was a problem hiding this comment.
Thanks for the explanation!
|
|
||
| The returned singleton instance uses the default MemoryPool. | ||
| """ | ||
| return MemoryManager.wrap(c_default_cpu_memory_manager()) |
There was a problem hiding this comment.
Do you plan to later subclass Device with CPUDevice and MemoryManager with CPUMemoryManager, respectively?
There was a problem hiding this comment.
Not necessarily, for now, if not needed? For CPU I am not sure if it is worth it, as looking at device.h, it seems the main thing the subclass provides is allowing to specify a memory pool when creating (which we could also expose through default_cpu_memory_manager if we want to expose this?
For CUDA there might be more CUDA-specific methods/attributes where it might be more worth providing a subclass.
There was a problem hiding this comment.
Indeed, we can do add them later when desired.
github-actions
bot
added
awaiting changes
We should certainly do that, but we were planning to do that in a follow-up |
Ah, fair enough. |
github-actions
bot
added
awaiting change review
Rationale for this change
Add bindings for the C++
arrow::Deviceandarrow::MemoryManagerclasses.What changes are included in this PR?
Basic bindings by adding the
pyarrow.Deviceandpyarrow.MemoryManagerclasses, and just tested for CPU.What is not included here are additional methods on the
MemoryManagerclass (eg to allocate or copy buffers), and this is also not yet tested for CUDA. Planning to do this as follow-ups, and first doing those basic bindings should enable further enhancements to be done in parallel.Are these changes tested?
Yes, for the CPU device only.