You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "jorisvandenbossche (via GitHub)" <gi...@apache.org> on 2024/03/21 16:26:26 UTC
[PR] GH-38325: [Python] Implement PyCapsule interface for Device data in PyArrow [arrow]
jorisvandenbossche opened a new pull request, #40717:
URL: https://github.com/apache/arrow/pull/40717
### Rationale for this change
PyArrow implementation for the specification additions being proposed in https://github.com/apache/arrow/pull/40708
### What changes are included in this PR?
New `__arrow_c_device_array__` method to `pyarrow.Array` and `pyarrow.RecordBatch`, and support in the `pyarrow.array(..)`, `pyarrow.record_batch(..)` and `pyarrow.table(..)` functions to consume objects that have those methods.
### Are these changes tested?
Yes (for CPU only for now, https://github.com/apache/arrow/pull/40385 is a prerequisite to test this for CUDA)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-38325: [Python] Implement PyCapsule interface for Device data in PyArrow [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on PR #40717:
URL: https://github.com/apache/arrow/pull/40717#issuecomment-2020889386
Thanks for testing!
> the only difference I see is that the device_id for the CPU is -1. I can't find in the spec exactly what it should be...is -1 the accepted value?
That's a good point. In practice this comes from our implementation in C++:
https://github.com/apache/arrow/blob/434f87274e8e9adab4f0434ae494f30dc955ca6e/cpp/src/arrow/device.h#L82-L86
But we should probably clarify that in the spec whether that's allowed / expected.
I see that DLPack actually specified that as to be 0 for CPU: https://dmlc.github.io/dlpack/latest/c_api.html#c.DLDevice.device_id
Maybe we should follow that; will open a separate issue.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-38325: [Python] Implement PyCapsule interface for Device data in PyArrow [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org