Copilot commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3932494990
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +261,232 @@ Result<DLDevice> ExportDevice(const
std::shared_ptr<Tensor>& t) {
return ExportDeviceImpl(t);
}
+/***************
+ * Consumers *
+ ***************/
+
+namespace {
+
+class CppDLTensor {
+ public:
+ using value_type = DLManagedTensorVersioned;
+ using pointer_type = value_type*;
+
+ static Result<CppDLTensor> TakeOwnership(pointer_type ptr) {
+ if (ARROW_PREDICT_FALSE(ptr == nullptr)) {
+ return Status::Invalid("Received null pointer.");
+ }
+ // Create the wrapper before checking the version as the spec mandates
that the
+ // deleter MUST be called on version major mismatch.
+ auto out = CppDLTensor(ptr);
+ if (ARROW_PREDICT_FALSE(out.ptr_->version.major != VERSION.major)) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", VERSION.major);
+ }
+ if (ARROW_PREDICT_FALSE(out.tensor().ndim < 0)) {
+ return Status::Invalid("Invalid DLPack tensor: ndim must be >= 0");
+ }
+ if (ARROW_PREDICT_FALSE(out.tensor().ndim != 0 && out.tensor().shape ==
nullptr)) {
+ return Status::Invalid(
+ "Invalid DLPack tensor: shape must be non-null when ndim != 0");
+ }
+ if (ARROW_PREDICT_FALSE(out.tensor().ndim != 0 && out.tensor().strides ==
nullptr)) {
+ return Status::Invalid(
+ "Invalid DLPack tensor: strides must be non-null when ndim != 0");
+ }
+ return out;
+ }
+
+ const DLTensor& tensor() const { return ptr_->dl_tensor; }
+
+ int64_t ndim() const {
+ DCHECK_GE(tensor().ndim, 0);
+ return tensor().ndim;
+ }
+
+ template <typename T>
+ T* data_as() {
+ return static_cast<T*>(tensor().data);
+ }
+
+ std::span<const int64_t> shape() const {
+ return {tensor().shape, static_cast<std::size_t>(ndim())};
+ }
+
+ std::span<const int64_t> strides() const {
+ return {tensor().strides, static_cast<std::size_t>(ndim())};
+ }
+
+ bool flag_is_set(uint8_t bits) const { return (ptr_->flags & bits) == bits; }
+
+ bool is_readonly() const { return
flag_is_set(DLPACK_FLAG_BITMASK_READ_ONLY); }
+
+ int32_t byte_width() const { return tensor().dtype.bits / 8; }
+
+ private:
+ struct Deleter {
+ void operator()(pointer_type ptr) {
+ // Null is valid in DLPack spec
+ if (auto del = ptr->deleter) {
+ del(ptr);
+ }
+ }
+ };
+
+ /// Make a safe wrapper that will delete the resource in case of exception.
+ std::unique_ptr<value_type, Deleter> ptr_;
+
+ explicit CppDLTensor(pointer_type ptr) : ptr_(ptr) {}
+};
+
+Result<std::shared_ptr<FixedWidthType>> DataTypeFromDLPack(DLDataType dtype) {
+ if (dtype.lanes != 1) {
+ return Status::TypeError("Only type with one lane are supported.");
+ }
+
+ auto constexpr as_fw = [](auto dt) {
+ return std::static_pointer_cast<FixedWidthType>(std::move(dt));
+ };
+
+ switch (dtype.code) {
+ case kDLInt: {
+ switch (dtype.bits) {
+ case 8:
+ return as_fw(int8());
+ case 16:
+ return as_fw(int16());
+ case 32:
+ return as_fw(int32());
+ case 64:
+ return as_fw(int64());
+ default:
+ return Status::Invalid("unsupported integer bit width ",
+ static_cast<int>(dtype.bits));
+ }
+ }
+ case kDLUInt: {
+ switch (dtype.bits) {
+ case 8:
+ return as_fw(uint8());
+ case 16:
+ return as_fw(uint16());
+ case 32:
+ return as_fw(uint32());
+ case 64:
+ return as_fw(uint64());
+ default:
+ return Status::Invalid("unsupported unsigned integer bit width ",
+ static_cast<int>(dtype.bits));
+ }
+ }
+ case kDLFloat: {
+ switch (dtype.bits) {
+ case 16:
+ return as_fw(float16());
+ case 32:
+ return as_fw(float32());
+ case 64:
+ return as_fw(float64());
+ default:
+ return Status::Invalid("unsupported float bit width ",
+ static_cast<int>(dtype.bits));
+ }
+ }
+ default: {
+ return Status::Invalid("unsupported DLPack type ",
static_cast<int>(dtype.code));
+ }
+ }
+}
+
+inline std::vector<int64_t> StridesInBytes(std::span<const int64_t> strides,
+ int64_t byte_width) {
+ std::vector<int64_t> out{};
+ out.reserve(strides.size());
+ for (const auto& s : strides) {
+ out.push_back(s * byte_width);
+ }
+ return out;
+}
Review Comment:
`StridesInBytes` multiplies `int64_t` strides by `byte_width` using `s *
byte_width`, which can overflow and is undefined behavior in C++. Since the
inputs come from an external DLPack capsule, this should use
`MultiplyWithOverflow` and fail (or force an error path) on overflow.
##########
python/pyarrow/array.pxi:
##########
@@ -2269,6 +2274,51 @@ cdef class Array(_PandasConvertible):
return pyarrow_wrap_array(array)
+ @staticmethod
+ def from_dlpack(x, /, *, device=None, copy=None):
+ """
+ Construct an Array from an object implementing the DLPack protocol.
+
+ Parameters
+ ----------
+ x : object
+ The input object containing array data, following the DLPack
+ protocol (has a ``__dlpack__`` method).
+ device : tuple[enum.Enum, int], optional
+ Designates where the resulting Array should reside, in the
+ format returned by :meth:`Array.__dlpack_device__`. When None,
+ the output Array occupies the same device as the source.
+ Default: None.
+ copy : bool, optional
+ Controls duplication behavior. True mandates copying; False
+ prohibits copying and raises ``BufferError`` if unavoidable;
+ None duplicates only when necessary. Default: None.
+
+ Returns
+ -------
+ Array
+ An Array housing the data from the input object, potentially
+ as a copy or view.
+ """
+ version = (DLPACK_VERSION.major, DLPACK_VERSION.minor)
+ pycapsule = x.__dlpack__(max_version=version, dl_device=device,
copy=copy)
+ if not PyCapsule_CheckExact(pycapsule):
+ raise TypeError("DLPack producer did not return a PyCapsule")
+ cdef DLManagedTensorVersioned* ptr =
<DLManagedTensorVersioned*>PyCapsule_GetPointer(
+ pycapsule, "dltensor_versioned")
+ if ptr == NULL:
+ raise ValueError(
+ 'DLPack producer did not produce a "dltensor_versioned"
PyCapsule')
+ # Mark the capsule as consumed so its destructor does not also invoke
the deleter.
+ # ImportArrayVersionedFromDLPack will take ownership even if it errors
(calling
+ # the deleter in that case).
+ PyCapsule_SetName(pycapsule, "used_dltensor_versioned")
+ with nogil:
+ # Copy handled on producer side
+ result = ImportArrayVersionedFromDLPack(ptr, False)
+ carray = GetResultValue(result)
+ return pyarrow_wrap_array(carray)
Review Comment:
`copy` is documented as “True mandates copying”, but the implementation
always calls `ImportArrayVersionedFromDLPack(ptr, False)`, so Arrow never
enforces an owning copy. This can violate the contract for producers that
ignore or only best-effort the `copy=` hint. Consider passing `copy is True`
down to the C++ importer so `copy=True` reliably produces an owning buffer.
##########
python/pyarrow/tensor.pxi:
##########
@@ -300,7 +306,52 @@ strides: {self.strides}"""
buffer.strides = <Py_ssize_t *>
cp.PyBytes_AsString(self._ssize_t_strides)
buffer.suboffsets = NULL
- def __dlpack__(self, stream=None, max_version=None, dl_device=None,
copy=None):
+ @staticmethod
+ def from_dlpack(x, /, *, device=None, copy=None):
+ """
+ Construct a Tensor from an object implementing the DLPack protocol.
+
+ Parameters
+ ----------
+ x : object
+ The input object containing array data, following the DLPack
+ protocol (has a ``__dlpack__`` method).
+ device : tuple[enum.Enum, int], optional
+ Designates where the resulting Tensor should reside, in the
+ format returned by :meth:`Tensor.__dlpack_device__`. When None,
+ the output Tensor occupies the same device as the source.
+ Default: None.
+ copy : bool, optional
+ Controls duplication behavior. True mandates copying; False
+ prohibits copying and raises ``BufferError`` if unavoidable;
+ None duplicates only when necessary. Default: None.
+
+ Returns
+ -------
+ Tensor
+ A Tensor housing the data from the input object, potentially
+ as a copy or view.
+ """
+ version = (DLPACK_VERSION.major, DLPACK_VERSION.minor)
+ pycapsule = x.__dlpack__(max_version=version, dl_device=device,
copy=copy)
+ if not PyCapsule_CheckExact(pycapsule):
+ raise TypeError("DLPack producer did not return a PyCapsule")
+ cdef DLManagedTensorVersioned* ptr =
<DLManagedTensorVersioned*>PyCapsule_GetPointer(
+ pycapsule, "dltensor_versioned")
+ if ptr == NULL:
+ raise ValueError(
+ 'DLPack producer did not produce a "dltensor_versioned"
PyCapsule')
+ # Mark the capsule as consumed so its destructor does not also invoke
the deleter.
+ # ImportTensorVersionedFromDLPack will take ownership even if it
errors (calling
+ # the deleter in that case).
+ PyCapsule_SetName(pycapsule, "used_dltensor_versioned")
+ with nogil:
+ # Copy handled on producer side
+ result = ImportTensorVersionedFromDLPack(ptr, False)
+ ctensor = GetResultValue(result)
+ return pyarrow_wrap_tensor(ctensor)
Review Comment:
`copy` is documented as “True mandates copying”, but the implementation
always calls `ImportTensorVersionedFromDLPack(ptr, False)`, so Arrow never
enforces an owning copy. This can violate the contract for producers that
ignore or only best-effort the `copy=` hint. Consider passing `copy is True`
down to the C++ importer so `copy=True` reliably produces an owning buffer.
--
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]