jorisvandenbossche commented on code in PR #40717:
URL: https://github.com/apache/arrow/pull/40717#discussion_r1534246954
##########
python/pyarrow/array.pxi:
##########
@@ -1855,6 +1868,74 @@ cdef class Array(_PandasConvertible):
)
return pyarrow_wrap_array(c_array)
+ def __arrow_c_device_array__(self, requested_schema=None):
+ """
+ Get a pair of PyCapsules containing a C ArrowDeviceArray representation
+ of the object.
+
+ Parameters
+ ----------
+ requested_schema : PyCapsule | None
+ A PyCapsule containing a C ArrowSchema representation of a
requested
+ schema. PyArrow will attempt to cast the array to this data type.
+ If None, the array will be returned as-is, with a type matching the
+ one returned by :meth:`__arrow_c_schema__()`.
+
+ Returns
+ -------
+ Tuple[PyCapsule, PyCapsule]
+ A pair of PyCapsules containing a C ArrowSchema and
ArrowDeviceArray,
+ respectively.
+ """
+ cdef:
+ ArrowDeviceArray* c_array
+ ArrowSchema* c_schema
+ shared_ptr[CArray] inner_array
+
+ if requested_schema is not None:
+ target_type = DataType._import_from_c_capsule(requested_schema)
+
+ if target_type != self.type:
+ # TODO should protect from trying to cast non-CPU data
+ try:
+ casted_array = _pc().cast(self, target_type, safe=True)
+ inner_array = pyarrow_unwrap_array(casted_array)
+ except ArrowInvalid as e:
+ raise ValueError(
+ f"Could not cast {self.type} to requested type
{target_type}: {e}"
+ )
+ else:
+ inner_array = self.sp_array
+ else:
+ inner_array = self.sp_array
Review Comment:
This part is a bit repetitive with the non-device version. I could factor
that out into a shared helper function
--
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]
