From 71dfc175f4c6a12f5b9dcb75f5c545f6057ea1df Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Tue, 23 Jan 2024 17:27:36 +0000 Subject: [PATCH] Expose the device interface low-level with Array._export_to_c_device method --- cpp/src/arrow/c/bridge.cc | 2 +- python/pyarrow/array.pxi | 27 +++++++++++++++++++++++++++ python/pyarrow/includes/libarrow.pxd | 11 +++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/c/bridge.cc b/cpp/src/arrow/c/bridge.cc index 238afb0328672..ea883a3a6738b 100644 --- a/cpp/src/arrow/c/bridge.cc +++ b/cpp/src/arrow/c/bridge.cc @@ -587,7 +587,7 @@ struct ArrayExporter { export_.buffers_.resize(n_buffers); std::transform(buffers_begin, data->buffers.end(), export_.buffers_.begin(), [](const std::shared_ptr& buffer) -> const void* { - return buffer ? buffer->data() : nullptr; + return buffer ? reinterpret_cast(buffer->address()) : nullptr; }); if (need_variadic_buffer_sizes) { diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 1416f5f4346d9..d043653335f9c 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -1682,6 +1682,33 @@ cdef class Array(_PandasConvertible): c_ptr, c_schema_ptr)) + def _export_to_c_device(self, out_ptr, out_schema_ptr): + """ + Export to a C ArrowDeviceArray struct, given its pointer. + + If a C ArrowSchema struct pointer is also given, the array type + is exported to it at the same time. + + Parameters + ---------- + out_ptr: int + The raw pointer to a C ArrowDeviceArray struct. + out_schema_ptr: int (optional) + The raw pointer to a C ArrowSchema struct. + + Be careful: if you don't pass the ArrowDeviceArray struct to a consumer, + array memory will leak. This is a low-level function intended for + expert users. + """ + cdef: + void* c_ptr = _as_c_pointer(out_ptr) + void* c_schema_ptr = _as_c_pointer(out_schema_ptr, + allow_null=True) + with nogil: + check_status(ExportDeviceArray( + deref(self.sp_array), NULL, + c_ptr, c_schema_ptr)) + @staticmethod def _import_from_c(in_ptr, type): """ diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 74e92594b04e5..acb260295258d 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -341,6 +341,12 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: CResult[unique_ptr[CResizableBuffer]] AllocateResizableBuffer( const int64_t size, CMemoryPool* pool) + cdef cppclass CSyncEvent" arrow::Device::SyncEvent": + pass + + cdef cppclass CDevice" arrow::Device": + pass + cdef CMemoryPool* c_default_memory_pool" arrow::default_memory_pool"() cdef CMemoryPool* c_system_memory_pool" arrow::system_memory_pool"() cdef CStatus c_jemalloc_memory_pool" arrow::jemalloc_memory_pool"( @@ -2802,6 +2808,9 @@ cdef extern from "arrow/c/abi.h": cdef struct ArrowArrayStream: void (*release)(ArrowArrayStream*) noexcept nogil + cdef struct ArrowDeviceArray: + pass + cdef extern from "arrow/c/bridge.h" namespace "arrow" nogil: CStatus ExportType(CDataType&, ArrowSchema* out) CResult[shared_ptr[CDataType]] ImportType(ArrowSchema*) @@ -2831,6 +2840,8 @@ cdef extern from "arrow/c/bridge.h" namespace "arrow" nogil: CResult[shared_ptr[CRecordBatchReader]] ImportRecordBatchReader( ArrowArrayStream*) + CStatus ExportDeviceArray(const CArray&, shared_ptr[CSyncEvent], + ArrowDeviceArray* out, ArrowSchema*) cdef extern from "arrow/util/byte_size.h" namespace "arrow::util" nogil: CResult[int64_t] ReferencedBufferSize(const CArray& array_data)