pitrou commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3927044410
##########
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`.
Review Comment:
A new `Buffer` subclass would be cleaner than this IMHO.
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -18,22 +18,37 @@
#include "arrow/c/dlpack.h"
#include <array>
+#include <functional>
#include <memory>
+#include <numeric>
#include <type_traits>
#include <vector>
#include "arrow/array/array_base.h"
+#include "arrow/array/util.h"
#include "arrow/buffer.h"
#include "arrow/c/dlpack_abi.h"
#include "arrow/device.h"
#include "arrow/tensor.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/logging_internal.h"
+#include "arrow/util/small_vector.h"
namespace arrow::dlpack {
+extern const DLPackVersion VERSION = {
Review Comment:
`kVersion` or `kDLPackVersion` rather than `VERSION`?
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +261,238 @@ 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));
+ }
+ }
+}
+
+Result<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) {
+ int64_t stride_bytes = 0;
+ if (ARROW_PREDICT_FALSE(
+ internal::MultiplyWithOverflow(s, byte_width, &stride_bytes))) {
+ return Status::Invalid("Overflow computing DLPack tensor stride in
bytes.");
+ }
+ out.push_back(stride_bytes);
+ }
+ 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));
+ int64_t nbytes = 0;
+ if (ARROW_PREDICT_FALSE(internal::MultiplyWithOverflow(
+ nelements, static_cast<int64_t>(dl.byte_width()), &nbytes))) {
+ return Status::Invalid("Overflow computing DLPack tensor size in bytes.");
+ }
+ // 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 {
Review Comment:
Presumably, a caller setting `copy` wants to get an exclusive mutable copy
of the tensor.
Should we then inspect the DLPack flags for
[`DLPACK_FLAG_BITMASK_IS_COPIED`](https://dmlc.github.io/dlpack/latest/c_api.html#c.DLPACK_FLAG_BITMASK_IS_COPIED)?
In that case, a new copy doesn't need to be made, I think?
(also check that `DLPACK_FLAG_BITMASK_READ_ONLY` is not set?)
##########
cpp/src/arrow/c/dlpack.h:
##########
@@ -105,4 +108,32 @@ Result<DLDevice> ExportDevice(const
std::shared_ptr<Array>& arr);
ARROW_EXPORT
Result<DLDevice> ExportDevice(const std::shared_ptr<Tensor>& t);
+/// \brief Import a DLPack tensor as an Arrow Array.
+///
+/// Same restrictions on data types as `ExportArrayVersioned`, only row-major
+/// tensors are supported. Takes ownership of the `DLManagedTensorVersioned` in
+/// an error-safe fashion.
+///
+/// \param[in] raw DLPack tensor
+/// \param[in] copy Whether to copy the data instead of sharing it with the
DLPack
+/// producer.
+/// \return An Arrow Array
+ARROW_EXPORT
+Result<std::shared_ptr<Array>> ImportArrayVersioned(DLManagedTensorVersioned*
raw,
Review Comment:
Is the "Versioned" suffix useful? We're not exposing any non-versioned API,
and even then, we could just rely on function overloading.
##########
python/pyarrow/tests/test_dlpack.py:
##########
@@ -339,3 +339,81 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_tensor_from_dlpack(np_type):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ expected = base[::2, 1::2]
+ assert not expected.flags['C_CONTIGUOUS']
+ tensor = pa.Tensor.from_dlpack(expected)
+ assert isinstance(tensor, pa.Tensor)
+ np.testing.assert_array_equal(tensor.to_numpy(), expected, strict=True)
+
+
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_array_from_dlpack(np_type):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
+ arr = pa.Array.from_dlpack(expected)
+ assert isinstance(arr, pa.Array)
+ np.testing.assert_array_equal(arr.to_numpy(), expected, strict=True)
+
+
+@check_bytes_allocated
+def test_from_dlpack_zero_copy():
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ expected = np.array([1, 2, 3], dtype=np.int64)
+ tensor = pa.Tensor.from_dlpack(expected)
+ result = tensor.to_numpy()
+ expected[0] = 100
+ # Zero-copy import: mutating the source is visible through the tensor.
+ assert result[0] == 100
Review Comment:
Can we mutate `result` too?
##########
python/pyarrow/tests/test_dlpack.py:
##########
@@ -339,3 +339,81 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_tensor_from_dlpack(np_type):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ expected = base[::2, 1::2]
+ assert not expected.flags['C_CONTIGUOUS']
+ tensor = pa.Tensor.from_dlpack(expected)
+ assert isinstance(tensor, pa.Tensor)
+ np.testing.assert_array_equal(tensor.to_numpy(), expected, strict=True)
Review Comment:
Should test that `tensor` is still valid after `expected` goes away?
##########
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`.
Review Comment:
That said, I understand that the subclassing strategy might not be scalable
once we want to support non-CPU devices.
##########
python/pyarrow/tests/test_dlpack.py:
##########
@@ -339,3 +339,81 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_tensor_from_dlpack(np_type):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ expected = base[::2, 1::2]
+ assert not expected.flags['C_CONTIGUOUS']
+ tensor = pa.Tensor.from_dlpack(expected)
+ assert isinstance(tensor, pa.Tensor)
+ np.testing.assert_array_equal(tensor.to_numpy(), expected, strict=True)
+
+
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_array_from_dlpack(np_type):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
+ arr = pa.Array.from_dlpack(expected)
Review Comment:
Call `arr.validate(full=True)`?
##########
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
Review Comment:
PyArrow also has its own Device class, can we add a TODO and/or open an
issue to support it here?
```python
>>> mm = pa.default_cpu_memory_manager()
>>> mm.device
<pyarrow.Device: CPUDevice()>
>>> mm.device.device_type
<DeviceAllocationType.CPU: 1>
>>> import pyarrow.cuda as cu
>>> ctx = cu.Context()
>>> ctx.device
<pyarrow.Device: CudaDevice(device_number=0, name="NVIDIA GeForce GT 1030")>
>>> ctx.device.device_type
<DeviceAllocationType.CUDA: 2>
```
##########
cpp/src/arrow/c/dlpack_test.cc:
##########
@@ -329,4 +329,353 @@ TYPED_TEST(TestExportTensor, TestTensorStrided) {
f_dlpack_strides);
}
+/***************
+ * Consumers *
+ ***************/
+
+/// A DLPack tensor as a foreign library would produce it.
+struct ForeignTensor {
+ DLDataType dtype = {.code = kDLFloat, .bits = 32, .lanes = 1};
+ std::vector<int64_t> shape = {};
+ /// In number of elements, as mandated by DLPack.
+ std::vector<int64_t> strides = {};
+ std::vector<uint8_t> data = {};
+ DLDevice device = {.device_type = kDLCPU, .device_id = 0};
+ uint64_t byte_offset = 0;
+ uint64_t flags = 0;
+ /// Incremented when the consumer releases the tensor.
+ std::shared_ptr<int> deleted = std::make_shared<int>(0);
+
+ DLManagedTensorVersioned managed = {};
+};
+
+template <typename T>
+std::vector<uint8_t> ToBytes(const std::vector<T>& values) {
+ std::vector<uint8_t> bytes(values.size() * sizeof(T));
+ std::memcpy(bytes.data(), values.data(), bytes.size());
+ return bytes;
+}
+
+/// Hand out a DLPack tensor owning ``foreign``, releasing it through its
deleter.
+DLManagedTensorVersioned* Produce(ForeignTensor foreign) {
+ auto owned = std::make_unique<ForeignTensor>(std::move(foreign));
+ owned->managed = {
+ .version = {.major = DLPACK_MAJOR_VERSION, .minor =
DLPACK_MINOR_VERSION},
+ .manager_ctx = owned.get(),
+ .deleter =
+ [](DLManagedTensorVersioned* self) {
+ auto* ctx = static_cast<ForeignTensor*>(self->manager_ctx);
+ ++(*ctx->deleted);
+ delete ctx;
+ },
+ .flags = owned->flags,
+ .dl_tensor =
+ {
+ .data = owned->data.data(),
+ .device = owned->device,
+ .ndim = static_cast<int32_t>(owned->shape.size()),
+ .dtype = owned->dtype,
+ .shape = owned->shape.data(),
+ .strides = owned->strides.data(),
+ .byte_offset = owned->byte_offset,
+ },
+ };
+ return &owned.release()->managed;
+}
+
+template <bool kCopy>
+struct TensorConsumer {
+ using Imported = std::shared_ptr<Tensor>;
+ static constexpr bool copy = kCopy;
+ static constexpr const char* name = copy ? "TensorCopied" : "TensorShared";
+
+ static Result<Imported> Import(DLManagedTensorVersioned* raw) {
+ return ImportTensorVersioned(raw, copy);
+ }
+ static std::shared_ptr<DataType> ValueType(const Imported& t) { return
t->type(); }
+ static const uint8_t* RawData(const Imported& t) { return t->raw_data(); }
+ static bool IsMutable(const Imported& t) { return t->is_mutable(); }
+ static int64_t Size(const Imported& t) { return t->size(); }
+};
Review Comment:
Can we add a validation method that would call `Tensor::Validate` and
`Array::ValidateFull`?
--
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]
