Copilot commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3932077157
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +263,217 @@ 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 (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 (out.ptr_->version.major != VERSION.major) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", VERSION.major);
+ }
+ return out;
+ }
Review Comment:
CppDLTensor::TakeOwnership() only DCHECKs that ndim is non-negative and
doesn't validate shape/strides pointers for ndim>0. In release builds a
malformed/older-version DLPack tensor (e.g., negative ndim or null strides) can
lead to huge spans and crashes when shape()/strides() are used.
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +263,217 @@ 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 (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 (out.ptr_->version.major != VERSION.major) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", VERSION.major);
+ }
+ 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;
+}
+
+Result<std::shared_ptr<Buffer>> ImportBuffer(CppDLTensor&& dl, bool copy) {
+ // DLPack strides are in number of elements, so is the size we compute from
them.
+ ARROW_ASSIGN_OR_RAISE(const auto nelements,
+ internal::ComputeTensorSize(dl.shape(), dl.strides(),
1));
+ const auto nbytes = nelements * dl.byte_width();
+ // DLPack mandates a null data pointer when the tensor holds no element, so
there is
+ // neither anything to share nor to copy.
Review Comment:
ImportBuffer computes `nbytes = nelements * dl.byte_width()` without an
overflow check. If `nelements` is large (but still fits in int64), `nbytes` can
overflow and lead to constructing a Buffer with an incorrect (wrapped) size,
which is unsafe.
##########
python/pyarrow/array.pxi:
##########
@@ -2269,6 +2274,49 @@ 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)
+ cdef DLManagedTensorVersioned* ptr =
<DLManagedTensorVersioned*>PyCapsule_GetPointer(
+ pycapsule, "dltensor_versioned")
Review Comment:
Array.from_dlpack calls PyCapsule_GetPointer() without first validating that
the producer returned a PyCapsule. If __dlpack__ returns a non-capsule object,
PyCapsule_GetPointer sets an error which is then overwritten by the custom
ValueError, producing a misleading exception type/message.
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +263,217 @@ 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 (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 (out.ptr_->version.major != VERSION.major) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", VERSION.major);
+ }
+ 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;
+}
+
+Result<std::shared_ptr<Buffer>> ImportBuffer(CppDLTensor&& dl, bool copy) {
+ // DLPack strides are in number of elements, so is the size we compute from
them.
+ ARROW_ASSIGN_OR_RAISE(const auto nelements,
+ internal::ComputeTensorSize(dl.shape(), dl.strides(),
1));
+ const auto nbytes = nelements * dl.byte_width();
+ // DLPack mandates a null data pointer when the tensor holds no element, so
there is
+ // neither anything to share nor to copy.
+ uint8_t* data =
+ (nbytes == 0) ? nullptr : dl.data_as<uint8_t>() +
dl.tensor().byte_offset;
+
+ std::shared_ptr<Buffer> buffer = nullptr;
+ if (nbytes == 0) {
+ // DLPack data pointer may be null on empty tensors
+ buffer = std::make_shared<Buffer>(data, nbytes);
+ } else if (copy) {
+ ARROW_ASSIGN_OR_RAISE(buffer, MutableBuffer::CopyNonOwned(
+ {data, nbytes},
default_cpu_memory_manager()));
+ } else {
+ const bool readonly = dl.is_readonly();
+ // Trick to keep DLPack data alive taken from `Buffer::FromVector`.
+ auto deleter = [dl = std::move(dl)](auto* buffer) { delete buffer; };
+ if (readonly) {
+ buffer = {new Buffer{data, nbytes}, std::move(deleter)};
+ } else {
+ buffer = std::shared_ptr<MutableBuffer>{
+ new MutableBuffer{data, nbytes},
+ std::move(deleter),
+ };
+ }
+ }
+
+ return buffer;
+}
+
+} // namespace
+
+Result<std::shared_ptr<Array>> ImportArrayVersioned(DLManagedTensorVersioned*
unmanaged,
+ bool copy) {
+ ARROW_ASSIGN_OR_RAISE(auto dl, CppDLTensor::TakeOwnership(unmanaged));
+
+ if (dl.tensor().device.device_type != kDLCPU) {
+ return Status::NotImplemented(
+ "DLPack support is implemented only for buffers on CPU device.");
+ }
+
+ if (dl.ndim() != 1 || dl.strides().front() != 1) {
+ return Status::Invalid(
+ "Only contiguous one dimensional tensor can be imported as arrays."
+ " Try importing to Tensor first.");
+ }
Review Comment:
ImportArrayVersioned returns Status::Invalid for unsupported ndim/stride
cases, but the tests (and Python API) expect a NotImplemented error for these
unsupported DLPack tensor layouts. Returning Invalid will surface as
ArrowInvalid/ValueError and will fail dlpack_test.cc expectations.
--
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]
