zeroshade commented on code in PR #39980:
URL: https://github.com/apache/arrow/pull/39980#discussion_r1496104359
##########
python/pyarrow/array.pxi:
##########
@@ -1778,6 +1778,70 @@ cdef class Array(_PandasConvertible):
return pyarrow_wrap_array(array)
+ def _export_to_c_device(self, out_ptr, out_schema_ptr=0):
+ """
+ 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), <shared_ptr[CSyncEvent]>NULL,
+ <ArrowDeviceArray*> c_ptr, <ArrowSchema*> c_schema_ptr))
+
+ @staticmethod
+ def _import_from_c_device(in_ptr, type):
+ """
+ Import Array from a C ArrowDeviceArray struct, given its pointer
+ and the imported array type.
+
+ Parameters
+ ----------
+ in_ptr: int
+ The raw pointer to a C ArrowDeviceArray struct.
+ type: DataType or int
+ Either a DataType object, or the raw pointer to a C ArrowSchema
+ struct.
+
+ This is a low-level function intended for expert users.
+ """
+ cdef:
+ void* c_ptr = _as_c_pointer(in_ptr)
+ void* c_type_ptr
+ shared_ptr[CArray] c_array
+
+ c_type = pyarrow_unwrap_data_type(type)
+ if c_type == nullptr:
+ # Not a DataType object, perhaps a raw ArrowSchema pointer
+ c_type_ptr = _as_c_pointer(type)
+ with nogil:
+ c_array = GetResultValue(
+ ImportDeviceArray(<ArrowDeviceArray*> c_ptr,
+ <ArrowSchema*> c_type_ptr)
+ )
Review Comment:
The default mapper is only allowing CPU arrays, but pyarrow does have a cuda
lib, shouldn't we allow and enable importing at least CUDA arrays too?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
