This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 39706ed0617 GH-39295: [C++][Python] ConsumingDLPack on Arrays and
Tensor (#51122)
39706ed0617 is described below
commit 39706ed0617da98351f98797048ca228c4f5c15e
Author: Antoine Prouvost <[email protected]>
AuthorDate: Thu Sep 10 17:04:57 2026 +0200
GH-39295: [C++][Python] ConsumingDLPack on Arrays and Tensor (#51122)
### Rationale for this change
There are already utilities to import tensor data from NumPy.
With DLPack standard, it will work across many tensor providers.
### What changes are included in this PR?
- Bind `FixedShapedTensorArray::FromTensor` in Python
- Add `ImportArrayVersionedFromDLPack` and
`ImportTensorVersionedFromDLPack` in C++
- Add in Python `Array.from_dlpack` and `Tensor.from_dlpack`
Similar to the export, multidimensional tensor import require an
explicit step through `Tensor`.
- No `Array::FromTensor` or even `FixedSizeListArray::FromTensor`. IMHO it
does not feel as necessary in this direction but open to anyone's take one it.
### Are these changes tested?
Yes.
### Are there any user-facing changes?
Yes, new APIs.
* GitHub Issue: #39295
Lead-authored-by: AntoinePrv <[email protected]>
Co-authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/buffer.cc | 19 +-
cpp/src/arrow/buffer.h | 67 ++++-
cpp/src/arrow/buffer_test.cc | 2 +-
cpp/src/arrow/c/dlpack.cc | 277 +++++++++++++++++++-
cpp/src/arrow/c/dlpack.h | 25 ++
cpp/src/arrow/c/dlpack_test.cc | 386 ++++++++++++++++++++++++++++
cpp/src/arrow/tensor.cc | 55 +++-
cpp/src/arrow/tensor.h | 9 +
docs/source/python/dlpack.rst | 124 ++++++++-
python/pyarrow/array.pxi | 116 ++++++++-
python/pyarrow/includes/libarrow.pxd | 20 ++
python/pyarrow/tensor.pxi | 53 +++-
python/pyarrow/tests/test_dlpack.py | 130 ++++++++++
python/pyarrow/tests/test_extension_type.py | 31 +++
14 files changed, 1259 insertions(+), 55 deletions(-)
diff --git a/cpp/src/arrow/buffer.cc b/cpp/src/arrow/buffer.cc
index 17e74520464..9abe41b4c3a 100644
--- a/cpp/src/arrow/buffer.cc
+++ b/cpp/src/arrow/buffer.cc
@@ -148,22 +148,13 @@ Result<std::shared_ptr<Buffer>> Buffer::ViewOrCopy(
return MemoryManager::CopyBuffer(source, to);
}
-class StlStringBuffer : public Buffer {
- public:
- explicit StlStringBuffer(std::string data) : input_(std::move(data)) {
- if (!input_.empty()) {
- data_ = reinterpret_cast<const uint8_t*>(input_.c_str());
- size_ = static_cast<int64_t>(input_.size());
- capacity_ = size_;
- }
+std::shared_ptr<Buffer> Buffer::FromString(std::string data) {
+ if (data.empty()) {
+ return std::shared_ptr<Buffer>{new Buffer()};
}
- private:
- std::string input_;
-};
-
-std::shared_ptr<Buffer> Buffer::FromString(std::string data) {
- return std::make_shared<StlStringBuffer>(std::move(data));
+ auto size_in_bytes = static_cast<int64_t>(data.size());
+ return TakeOwnership(std::move(data), size_in_bytes);
}
std::shared_ptr<Buffer> SliceMutableBuffer(std::shared_ptr<Buffer> buffer,
diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h
index 07f9931eba7..388829eba05 100644
--- a/cpp/src/arrow/buffer.h
+++ b/cpp/src/arrow/buffer.h
@@ -24,6 +24,7 @@
#include <span>
#include <string>
#include <string_view>
+#include <type_traits>
#include <utility>
#include <vector>
@@ -50,6 +51,15 @@ namespace arrow {
///
/// The following invariant is always true: Size <= Capacity
class ARROW_EXPORT Buffer {
+ private:
+ /// \brief Default data accessor used by TakeOwnership.
+ struct DefaultGetData {
+ template <typename T>
+ auto* operator()(T& container) const {
+ return container.data();
+ }
+ };
+
public:
ARROW_DISALLOW_COPY_AND_ASSIGN(Buffer);
@@ -146,14 +156,58 @@ class ARROW_EXPORT Buffer {
}
}
- /// \brief Construct an immutable buffer that takes ownership of the contents
+ /// \brief Construct a buffer that takes ownership of a container.
+ ///
+ /// This operation does not make a copy. If the underlying container is
mutable (as
+ /// detected by the return type of `get_data`) then returned buffer will be
mutable.
+ ///
+ /// \param[in] container The container to own. The container must own its
data as a
+ /// contiguous slice. That data does not need to remain at a
stable address
+ /// across a container move.
+ /// \param[in] nbytes The size of the data, which must not exceed the number
of bytes
+ /// readable from the pointer returned by \p get_data
+ /// \param[in] get_data Callable returning the address of the container's
data. This
+ /// callable is invoked *after* the container has been moved to
its final
+ /// address, to work with types such as `std::string`.
+ /// \return a new Buffer instance
+ template <typename T, typename Func = DefaultGetData>
+ static auto TakeOwnership(T container, int64_t nbytes, Func&& get_data = {})
{
+ using DataPtr = decltype(std::forward<Func>(get_data)(container));
+ constexpr bool kIsMutable =
!std::is_const_v<std::remove_pointer_t<DataPtr>>;
+ using BufferType = std::conditional_t<kIsMutable, MutableBuffer, Buffer>;
+ using Byte = std::conditional_t<kIsMutable, uint8_t, const uint8_t>;
+
+ // Hold the container and the Buffer in a single allocation. Declaration
order
+ // matters: the container is constructed first and destroyed last, so the
Buffer
+ // never outlives the memory it points into.
+ struct ControlBlock {
+ T container;
+ BufferType buffer;
+
+ ControlBlock(T container, int64_t nbytes, Func&& get_data)
+ : container(std::move(container)),
+ // Read the data pointer only once the container has reached its
final
+ // address, since moving it may invalidate the pointer (e.g. in a
small
+ // string optimization).
+
buffer(reinterpret_cast<Byte*>(std::forward<Func>(get_data)(this->container)),
+ nbytes) {}
+ };
+
+ auto owner = std::make_shared<ControlBlock>(std::move(container), nbytes,
+ std::forward<Func>(get_data));
+ // Aliasing constructor
+ auto* buffer = &owner->buffer;
+ return std::shared_ptr<BufferType>{std::move(owner), buffer};
+ }
+
+ /// \brief Construct a mutable buffer that takes ownership of the contents
/// of an std::string (without copying it).
///
/// \param[in] data a string to own
/// \return a new Buffer instance
static std::shared_ptr<Buffer> FromString(std::string data);
- /// \brief Construct an immutable buffer that takes ownership of the contents
+ /// \brief Construct a mutable buffer that takes ownership of the contents
/// of an std::vector (without copying it). Only vectors of TrivialType
objects
/// (integers, floating point numbers, ...) can be wrapped by this function.
///
@@ -168,15 +222,8 @@ class ARROW_EXPORT Buffer {
return std::shared_ptr<Buffer>{new Buffer()};
}
- auto* data = reinterpret_cast<uint8_t*>(vec.data());
auto size_in_bytes = static_cast<int64_t>(vec.size() * sizeof(T));
- return std::shared_ptr<Buffer>{
- new Buffer{data, size_in_bytes},
- // Keep the vector's buffer alive inside the shared_ptr's destructor
until after
- // we have deleted the Buffer. Note we can't use this trick in
FromString since
- // std::string's data is inline for short strings so moving
invalidates pointers
- // into the string's buffer.
- [vec = std::move(vec)](Buffer* buffer) { delete buffer; }};
+ return TakeOwnership(std::move(vec), size_in_bytes);
}
/// \brief Create buffer referencing typed memory with some length without
diff --git a/cpp/src/arrow/buffer_test.cc b/cpp/src/arrow/buffer_test.cc
index 4dd210076ed..bdda0b4a747 100644
--- a/cpp/src/arrow/buffer_test.cc
+++ b/cpp/src/arrow/buffer_test.cc
@@ -586,7 +586,7 @@ TEST(TestBuffer, FromStringRvalue) {
AssertIsCPUBuffer(*buffer);
}
- ASSERT_FALSE(buffer->is_mutable());
+ ASSERT_TRUE(buffer->is_mutable());
ASSERT_EQ(0, memcmp(buffer->data(), expected.c_str(), expected.size()));
ASSERT_EQ(static_cast<int64_t>(expected.size()), buffer->size());
diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc
index 4e25d50bb56..1b58b5a503c 100644
--- a/cpp/src/arrow/c/dlpack.cc
+++ b/cpp/src/arrow/c/dlpack.cc
@@ -19,21 +19,35 @@
#include <array>
#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/int_util_overflow.h"
+#include "arrow/util/logging_internal.h"
+#include "arrow/util/macros.h"
namespace arrow::dlpack {
+extern const DLPackVersion kVersion = {
+ .major = DLPACK_MAJOR_VERSION,
+ .minor = DLPACK_MINOR_VERSION,
+};
+
namespace {
+/***************
+ * Producers *
+ ***************/
+
Result<DLDataType> GetDLDataType(const DataType& type) {
auto dtype = DLDataType{};
dtype.lanes = 1;
@@ -108,7 +122,7 @@ DT* ExportBuffer(ExportBufferParams<Vec>&& p) {
// Strides must be non-null when ndim > 0
ctx->tensor.dl_tensor.strides = ctx->strides.data();
if constexpr (std::is_same_v<DT, DLManagedTensorVersioned>) {
- ctx->tensor.version = {.major = DLPACK_MAJOR_VERSION, .minor =
DLPACK_MINOR_VERSION};
+ ctx->tensor.version = kVersion;
ctx->tensor.flags = p.flags;
}
@@ -248,4 +262,265 @@ 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 != kVersion.major)) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", kVersion.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");
+ }
+ // Null strides are handled as row 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())};
+ }
+
+ /// Strides or empty span for old DLPack row-major convention.
+ std::span<const int64_t> strides() const {
+ if (auto strides = tensor().strides; strides != nullptr) {
+ return {strides, static_cast<std::size_t>(ndim())};
+ }
+ return {};
+ }
+
+ 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; }
+
+ /// Number of element in this tensor's buffer.
+ ///
+ /// Possibly more elements than represented in the tensor for non-contiguous
tensors.
+ ///
+ /// A zero dimensional tensor is a scalar, it holds a single element.
+ Result<int64_t> ComputeNumElements() const {
+ const auto strides = this->strides();
+ const auto shape = this->shape();
+ if (strides.size() > 0) {
+ // DLPack strides are in number of elements, so is the size we compute
from them.
+ return internal::ComputeTensorSize(shape, strides, 1);
+ }
+ // DLPack <1.3 may set strides == nullptr for row major
+ return std::reduce(shape.begin(), shape.end(), int64_t{1},
std::multiplies{});
+ }
+
+ /// Number of bytes needed to store this tensor data.
+ Result<int64_t> ComputeNumBytes() const {
+ ARROW_ASSIGN_OR_RAISE(const auto nelements, ComputeNumElements());
+ int64_t nbytes = 0;
+ if (ARROW_PREDICT_FALSE(internal::MultiplyWithOverflow(
+ nelements, static_cast<int64_t>(byte_width()), &nbytes))) {
+ return Status::Invalid("Overflow computing DLPack tensor size in
bytes.");
+ }
+ return nbytes;
+ }
+
+ 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) {
+ ARROW_ASSIGN_OR_RAISE(const int64_t nbytes, dl.ComputeNumBytes());
+
+ // 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 (const auto byte_offset = dl.tensor().byte_offset;
dl.is_readonly()) {
+ auto get_data = [byte_offset](auto& d) {
+ return d.template data_as<const uint8_t>() + byte_offset;
+ };
+ buffer = Buffer::TakeOwnership(std::move(dl), nbytes, get_data);
+ ARROW_DCHECK(!buffer->is_mutable());
+ } else {
+ auto get_data = [byte_offset](auto& d) {
+ return d.template data_as<uint8_t>() + byte_offset;
+ };
+ buffer = Buffer::TakeOwnership(std::move(dl), nbytes, get_data);
+ ARROW_DCHECK(buffer->is_mutable());
+ }
+
+ return buffer;
+}
+
+} // namespace
+
+Result<std::shared_ptr<Array>> ImportArrayVersioned(DLManagedTensorVersioned*
unmanaged) {
+ 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.");
+ }
+
+ const auto strides = dl.strides();
+ if (dl.ndim() != 1 || (!strides.empty() && strides.front() != 1)) {
+ return Status::Invalid(
+ "Only contiguous one dimensional tensor can be imported as arrays."
+ " Try importing to Tensor first.");
+ }
+
+ ARROW_ASSIGN_OR_RAISE(auto type, DataTypeFromDLPack(dl.tensor().dtype));
+ const auto nelements = dl.shape().front();
+ ARROW_ASSIGN_OR_RAISE(auto buffer, ImportBuffer(std::move(dl)));
+ auto data = ArrayData::Make(type, nelements, {nullptr, std::move(buffer)});
+ return MakeArray(std::move(data));
+}
+
+Result<std::shared_ptr<Tensor>> ImportTensorVersioned(
+ DLManagedTensorVersioned* unmanaged) {
+ 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.");
+ }
+
+ ARROW_ASSIGN_OR_RAISE(auto type, DataTypeFromDLPack(dl.tensor().dtype));
+ auto shape = std::vector<int64_t>(dl.shape().begin(), dl.shape().end());
+ const auto strides = dl.strides();
+
+ ARROW_ASSIGN_OR_RAISE(auto buffer, ImportBuffer(std::move(dl)));
+ const auto byte_width = type->byte_width();
+
+ // In older DLPack null strides means row major, same as in Arrow.
+ if (strides.empty()) {
+ return Tensor::Make(std::move(type), std::move(buffer), std::move(shape));
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto strides_bytes,
+ StridesInBytes(std::vector<int64_t>(strides.begin(), strides.end()),
byte_width));
+ return Tensor::Make(std::move(type), std::move(buffer), std::move(shape),
+ std::move(strides_bytes));
+}
+
} // namespace arrow::dlpack
diff --git a/cpp/src/arrow/c/dlpack.h b/cpp/src/arrow/c/dlpack.h
index 8a9084f36c7..29738dff6d1 100644
--- a/cpp/src/arrow/c/dlpack.h
+++ b/cpp/src/arrow/c/dlpack.h
@@ -25,6 +25,9 @@
namespace arrow::dlpack {
+/// The DLPack version used during compilation.
+ARROW_EXPORT extern const DLPackVersion kVersion;
+
/// \brief Export Arrow array as DLPack tensor.
///
/// DLMangedTensor is produced as defined by the DLPack protocol,
@@ -105,4 +108,26 @@ 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
+/// \return An Arrow Array
+ARROW_EXPORT
+Result<std::shared_ptr<Array>> ImportArrayVersioned(DLManagedTensorVersioned*
raw);
+
+/// \brief Import a DLPack tensor as an Arrow Tensor.
+///
+/// Same restrictions on data types as `ExportTensorVersioned`.
+/// Takes ownership of the `DLManagedTensorVersioned` in an error-safe fashion.
+/// If the DLPack input is marked as readonly, this will produce an immutable
tensor.
+///
+/// \param[in] raw DLPack tensor
+/// \return An Arrow Tensor
+ARROW_EXPORT
+Result<std::shared_ptr<Tensor>>
ImportTensorVersioned(DLManagedTensorVersioned* raw);
+
} // namespace arrow::dlpack
diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc
index 05de22237a9..8cbff0692da 100644
--- a/cpp/src/arrow/c/dlpack_test.cc
+++ b/cpp/src/arrow/c/dlpack_test.cc
@@ -329,4 +329,390 @@ 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;
+}
+
+struct TensorConsumer {
+ using Imported = std::shared_ptr<Tensor>;
+ static constexpr const char* name = "Tensor";
+
+ static Result<Imported> Import(DLManagedTensorVersioned* raw) {
+ return ImportTensorVersioned(raw);
+ }
+ 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(); }
+ static Status Validate(const Imported& t) { return t->Validate(); }
+
+ /// Import and additionally check the result is a well-formed Arrow object.
+ static Result<Imported> ImportAndValidate(DLManagedTensorVersioned* raw) {
+ ARROW_ASSIGN_OR_RAISE(auto imported, Import(raw));
+ RETURN_NOT_OK(Validate(imported));
+ return imported;
+ }
+};
+
+struct ArrayConsumer {
+ using Imported = std::shared_ptr<Array>;
+ static constexpr const char* name = "Array";
+
+ static Result<Imported> Import(DLManagedTensorVersioned* raw) {
+ return ImportArrayVersioned(raw);
+ }
+ static std::shared_ptr<DataType> ValueType(const Imported& arr) { return
arr->type(); }
+ static const uint8_t* RawData(const Imported& arr) {
+ return arr->data()->buffers[1]->data() + arr->offset() *
arr->type()->byte_width();
+ }
+ static bool IsMutable(const Imported& arr) {
+ return arr->data()->buffers[1]->is_mutable();
+ }
+ static int64_t Size(const Imported& arr) { return arr->length(); }
+ static Status Validate(const Imported& t) { return t->ValidateFull(); }
+
+ /// Import and additionally check the result is a well-formed Arrow object.
+ static Result<Imported> ImportAndValidate(DLManagedTensorVersioned* raw) {
+ ARROW_ASSIGN_OR_RAISE(auto imported, Import(raw));
+ RETURN_NOT_OK(Validate(imported));
+ return imported;
+ }
+};
+
+struct ConsumerNames {
+ template <typename Consumer>
+ static std::string GetName(int) {
+ return Consumer::name;
+ }
+};
+
+using ConsumerTypes = ::testing::Types<TensorConsumer, ArrayConsumer>;
+using TensorConsumerTypes = ::testing::Types<TensorConsumer>;
+using ArrayConsumerTypes = ::testing::Types<ArrayConsumer>;
+
+/// Tests sharing the same expectations for Arrow Tensor and Array imports.
+template <typename Consumer>
+class TestImport : public ::testing::Test {};
+
+TYPED_TEST_SUITE(TestImport, ConsumerTypes, ConsumerNames);
+
+TYPED_TEST(TestImport, Basic) {
+ auto foreign = ForeignTensor{
+ .shape = {6},
+ .strides = {1},
+ .data = ToBytes(std::vector<float>{0, 0, 1, 2, 3, 4, 5, 6}),
+ .byte_offset = 2 * sizeof(float),
+ .flags = DLPACK_FLAG_BITMASK_READ_ONLY,
+ };
+ const auto deleted = foreign.deleted;
+ const auto expected = std::vector<float>{1, 2, 3, 4, 5, 6};
+ const auto* values = foreign.data.data() + foreign.byte_offset;
+
+ ASSERT_OK_AND_ASSIGN(auto imported,
+
TypeParam::ImportAndValidate(Produce(std::move(foreign))));
+
+ AssertTypeEqual(*float32(), *TypeParam::ValueType(imported));
+ ASSERT_EQ(6, TypeParam::Size(imported));
+ ASSERT_FALSE(TypeParam::IsMutable(imported));
+ ASSERT_EQ(0, std::memcmp(TypeParam::RawData(imported), expected.data(),
+ expected.size() * sizeof(float)));
+
+ ASSERT_EQ(TypeParam::RawData(imported), values);
+ // The producer tensor is kept alive by the imported data
+ ASSERT_EQ(0, *deleted);
+ imported.reset();
+ ASSERT_EQ(1, *deleted);
+}
+
+TYPED_TEST(TestImport, Mutable) {
+ auto foreign = ForeignTensor{
+ .shape = {4},
+ .strides = {1},
+ .data = ToBytes(std::vector<float>{1, 2, 3, 4}),
+ .flags = 0,
+ };
+ ASSERT_OK_AND_ASSIGN(auto imported,
+
TypeParam::ImportAndValidate(Produce(std::move(foreign))));
+ ASSERT_TRUE(TypeParam::IsMutable(imported));
+}
+
+TYPED_TEST(TestImport, NullDeleter) {
+ // The DLPack spec allows producers not to set a deleter
+ auto* managed =
+ Produce({.shape = {2}, .strides = {1}, .data = std::vector<uint8_t>(8)});
+ auto* foreign = static_cast<ForeignTensor*>(managed->manager_ctx);
+ managed->deleter = nullptr;
+
+ ASSERT_OK_AND_ASSIGN(auto imported, TypeParam::ImportAndValidate(managed));
+ imported.reset();
+ delete foreign;
+}
+
+TYPED_TEST(TestImport, DataTypes) {
+ const std::vector<std::pair<DLDataType, std::shared_ptr<DataType>>> cases = {
+ {{kDLInt, 8, 1}, int8()}, {{kDLInt, 16, 1}, int16()},
+ {{kDLInt, 32, 1}, int32()}, {{kDLInt, 64, 1}, int64()},
+ {{kDLUInt, 8, 1}, uint8()}, {{kDLUInt, 16, 1}, uint16()},
+ {{kDLUInt, 32, 1}, uint32()}, {{kDLUInt, 64, 1}, uint64()},
+ {{kDLFloat, 16, 1}, float16()}, {{kDLFloat, 32, 1}, float32()},
+ {{kDLFloat, 64, 1}, float64()}};
+
+ for (const auto& [dtype, expected] : cases) {
+ ARROW_SCOPED_TRACE("dtype ", expected->ToString());
+ ASSERT_OK_AND_ASSIGN(
+ auto imported, TypeParam::ImportAndValidate(Produce(
+ {.dtype = dtype,
+ .shape = {3},
+ .strides = {1},
+ .data = std::vector<uint8_t>(3 *
expected->byte_width())})));
+ AssertTypeEqual(*expected, *TypeParam::ValueType(imported));
+ }
+}
+
+TYPED_TEST(TestImport, Empty) {
+ // DLPack mandates a null data pointer when the tensor holds no element
+ auto* managed = Produce({.shape = {0}, .strides = {1}});
+ managed->dl_tensor.data = nullptr;
+
+ ASSERT_OK_AND_ASSIGN(auto imported, TypeParam::ImportAndValidate(managed));
+ ASSERT_EQ(0, TypeParam::Size(imported));
+}
+
+TYPED_TEST(TestImport, Errors) {
+ auto check = [](ForeignTensor foreign, const std::string& message) {
+ const auto deleted = foreign.deleted;
+ const auto status =
+ TypeParam::ImportAndValidate(Produce(std::move(foreign))).status();
+ EXPECT_EQ(message, status.ToStringWithoutContextLines());
+ // Ownership is taken even when the import fails
+ EXPECT_EQ(1, *deleted);
+ };
+
+ ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Received null pointer.",
+ TypeParam::ImportAndValidate(nullptr));
+ check({.shape = {2},
+ .strides = {1},
+ .data = std::vector<uint8_t>(8),
+ .device = {.device_type = kDLCUDA, .device_id = 0}},
+ "NotImplemented: DLPack support is implemented only for buffers on CPU
device.");
+ check({.dtype = {kDLFloat, 32, 2}, .shape = {2}, .strides = {1}},
+ "Type error: Only type with one lane are supported.");
+ check({.dtype = {kDLInt, 4, 1}, .shape = {2}, .strides = {1}},
+ "Invalid: unsupported integer bit width 4");
+ check({.dtype = {kDLBool, 8, 1}, .shape = {2}, .strides = {1}},
+ "Invalid: unsupported DLPack type " + std::to_string(kDLBool));
+}
+
+TYPED_TEST(TestImport, UnsupportedVersion) {
+ auto* managed =
+ Produce({.shape = {2}, .strides = {1}, .data = std::vector<uint8_t>(8)});
+ const auto deleted =
static_cast<ForeignTensor*>(managed->manager_ctx)->deleted;
+ const auto major = DLPACK_MAJOR_VERSION + 1;
+ managed->version.major = major;
+
+ ASSERT_RAISES_WITH_MESSAGE(Invalid,
+ "Invalid: Unsupported DLPack major version " +
+ std::to_string(major) + ", expected " +
+ std::to_string(DLPACK_MAJOR_VERSION),
+ TypeParam::ImportAndValidate(managed));
+ // The spec mandates the deleter to be called on major version mismatch
+ ASSERT_EQ(1, *deleted);
+}
+
+template <typename Consumer>
+class TestImportTensor : public ::testing::Test {};
+
+TYPED_TEST_SUITE(TestImportTensor, TensorConsumerTypes, ConsumerNames);
+
+TYPED_TEST(TestImportTensor, ShapeAndStrides) {
+ auto foreign = ForeignTensor{
+ .shape = {2, 3},
+ .strides = {3, 1},
+ .data = ToBytes(std::vector<float>{1, 2, 3, 4, 5, 6}),
+ };
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+
TypeParam::ImportAndValidate(Produce(std::move(foreign))));
+
+ ASSERT_THAT(tensor->shape(), ::testing::ElementsAre(2, 3));
+ // Arrow strides are in bytes, DLPack strides in elements
+ ASSERT_THAT(tensor->strides(),
+ ::testing::ElementsAre(3 * sizeof(float), sizeof(float)));
+}
+
+TYPED_TEST(TestImportTensor, Empty) {
+ auto* managed = Produce({.shape = {0, 3}, .strides = {3, 1}});
+ managed->dl_tensor.data = nullptr;
+
+ ASSERT_OK_AND_ASSIGN(auto tensor, TypeParam::ImportAndValidate(managed));
+ ASSERT_THAT(tensor->shape(), ::testing::ElementsAre(0, 3));
+}
+
+TYPED_TEST(TestImportTensor, Strided) {
+ auto column_major = ForeignTensor{
+ .shape = {2, 3},
+ .strides = {1, 2},
+ .data = ToBytes(std::vector<float>{1, 2, 3, 4, 5, 6}),
+ };
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+
TypeParam::ImportAndValidate(Produce(std::move(column_major))));
+ ASSERT_TRUE(tensor->is_column_major());
+
+ // A 2x2 window over every other row of a 4x2 buffer
+ auto non_contiguous = ForeignTensor{
+ .shape = {2, 2},
+ .strides = {4, 1},
+ .data = ToBytes(std::vector<float>{1, 2, 3, 4, 5, 6, 7, 8}),
+ };
+ ASSERT_OK_AND_ASSIGN(tensor,
+
TypeParam::ImportAndValidate(Produce(std::move(non_contiguous))));
+ ASSERT_FALSE(tensor->is_contiguous());
+ ASSERT_EQ(6, tensor->template Value<FloatType>({1, 1}));
+}
+
+TYPED_TEST(TestImportTensor, ZeroDimensionIsScalar) {
+ auto* managed = Produce({
+ .shape = {},
+ .data = ToBytes(std::vector<float>{42}),
+ });
+ managed->dl_tensor.shape = nullptr;
+ managed->dl_tensor.strides = nullptr;
+
+ ASSERT_OK_AND_ASSIGN(auto tensor, TypeParam::ImportAndValidate(managed));
+ ASSERT_THAT(tensor->shape(), ::testing::IsEmpty());
+ ASSERT_EQ(1, tensor->size());
+ ASSERT_EQ(sizeof(float), tensor->data()->size());
+ ASSERT_EQ(42, tensor->template Value<FloatType>({}));
+}
+
+TYPED_TEST(TestImportTensor, NullStrides) {
+ // DLPack < 1.3 uses null strides to mean row major
+ auto* managed = Produce({
+ .shape = {2, 3},
+ .data = ToBytes(std::vector<float>{1, 2, 3, 4, 5, 6}),
+ });
+ managed->dl_tensor.strides = nullptr;
+
+ ASSERT_OK_AND_ASSIGN(auto tensor, TypeParam::ImportAndValidate(managed));
+ ASSERT_THAT(tensor->shape(), ::testing::ElementsAre(2, 3));
+ ASSERT_TRUE(tensor->is_row_major());
+ ASSERT_EQ(6, tensor->template Value<FloatType>({1, 2}));
+}
+
+TYPED_TEST(TestImportTensor, NegativeStrides) {
+ auto foreign = ForeignTensor{
+ .shape = {2, 2},
+ .strides = {-2, 1},
+ .data = std::vector<uint8_t>(16),
+ };
+ ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: negative strides not
supported",
+
TypeParam::ImportAndValidate(Produce(std::move(foreign))));
+}
+
+TYPED_TEST(TestImportTensor, RoundTrip) {
+ const auto original = TensorFromJSON(float64(), "[1, 2, 3, 4, 5, 6]", {3,
2});
+
+ ASSERT_OK_AND_ASSIGN(auto* managed, ExportTensorVersioned(original,
/*copy=*/false));
+ ASSERT_OK_AND_ASSIGN(auto tensor, TypeParam::ImportAndValidate(managed));
+
+ ASSERT_TRUE(tensor->Equals(*original));
+ ASSERT_EQ(original->raw_data(), tensor->raw_data());
+}
+
+template <typename Consumer>
+class TestImportArray : public ::testing::Test {};
+
+TYPED_TEST_SUITE(TestImportArray, ArrayConsumerTypes, ConsumerNames);
+
+TYPED_TEST(TestImportArray, OneDimension) {
+ auto foreign = ForeignTensor{
+ .dtype = {.code = kDLInt, .bits = 32, .lanes = 1},
+ .shape = {4},
+ .strides = {1},
+ .data = ToBytes(std::vector<int32_t>{1, 2, 3, 4}),
+ };
+ ASSERT_OK_AND_ASSIGN(auto array,
+
TypeParam::ImportAndValidate(Produce(std::move(foreign))));
+ AssertArraysEqual(*ArrayFromJSON(int32(), "[1, 2, 3, 4]"), *array);
+}
+
+TYPED_TEST(TestImportArray, Unsupported) {
+ auto check = [](ForeignTensor foreign) {
+ ASSERT_RAISES_WITH_MESSAGE(
+ Invalid,
+ "Invalid: Only contiguous one dimensional tensor can be imported as"
+ " arrays. Try importing to Tensor first.",
+ TypeParam::ImportAndValidate(Produce(std::move(foreign))));
+ };
+
+ // Only a Tensor can hold more than one dimension
+ check({.shape = {2, 3},
+ .strides = {3, 1},
+ .data = ToBytes(std::vector<float>{1, 2, 3, 4, 5, 6})});
+ check({.shape = {0, 3}, .strides = {1, 1}});
+ // Array values are contiguous, whatever the dimension count
+ check(
+ {.shape = {3}, .strides = {2}, .data = ToBytes(std::vector<float>{1, 2,
3, 4, 5})});
+ check({.shape = {2, 2}, .strides = {-2, 1}, .data =
std::vector<uint8_t>(16)});
+}
+
+TYPED_TEST(TestImportArray, RoundTrip) {
+ const auto original = ArrayFromJSON(float64(), "[1, 2, 3, 4, 5, 6]");
+
+ ASSERT_OK_AND_ASSIGN(auto* managed, ExportArrayVersioned(original,
/*copy=*/false));
+ ASSERT_OK_AND_ASSIGN(auto array, TypeParam::ImportAndValidate(managed));
+
+ AssertArraysEqual(*original, *array);
+ ASSERT_EQ(TypeParam::RawData(original), TypeParam::RawData(array));
+}
+
} // namespace arrow::dlpack
diff --git a/cpp/src/arrow/tensor.cc b/cpp/src/arrow/tensor.cc
index f2ff11a4f66..8b3137d66a3 100644
--- a/cpp/src/arrow/tensor.cc
+++ b/cpp/src/arrow/tensor.cc
@@ -109,13 +109,40 @@ Status ComputeColumnMajorStrides(const FixedWidthType&
type,
return Status::OK();
}
-} // namespace internal
+Result<int64_t> ComputeTensorSize(std::span<const int64_t> shape,
+ std::span<const int64_t> strides, int64_t
elem_size) {
+ // Check the largest offset can be computed without overflow
+ const size_t ndim = shape.size();
+ int64_t largest_offset = elem_size;
+ for (size_t i = 0; i < ndim; ++i) {
+ if (shape[i] == 0) continue;
+ if (strides[i] < 0) {
+ // TODO(mrkn): Support negative strides for sharing views
+ return Status::Invalid("negative strides not supported");
+ }
-namespace {
+ int64_t dim_offset = 0;
+ if (!internal::MultiplyWithOverflow(shape[i] - 1, strides[i],
&dim_offset)) {
+ if (!internal::AddWithOverflow(largest_offset, dim_offset,
&largest_offset)) {
+ continue;
+ }
+ }
-inline bool IsTensorStridesRowMajor(const std::shared_ptr<DataType>& type,
- const std::vector<int64_t>& shape,
- const std::vector<int64_t>& strides) {
+ return Status::Invalid(
+ "offsets computed from shape and strides would not fit in 64-bit
integer");
+ }
+
+ // A dimension with no element means empty for which the preceding does not
apply.
+ if (std::find(shape.begin(), shape.end(), 0) != shape.end()) {
+ return 0;
+ }
+ return largest_offset;
+}
+
+namespace {
+bool IsTensorStridesRowMajor(const std::shared_ptr<DataType>& type,
+ const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& strides) {
std::vector<int64_t> c_strides;
const auto& fw_type = checked_cast<const FixedWidthType&>(*type);
if (internal::ComputeRowMajorStrides(fw_type, shape, &c_strides).ok()) {
@@ -125,9 +152,9 @@ inline bool IsTensorStridesRowMajor(const
std::shared_ptr<DataType>& type,
}
}
-inline bool IsTensorStridesColumnMajor(const std::shared_ptr<DataType>& type,
- const std::vector<int64_t>& shape,
- const std::vector<int64_t>& strides) {
+bool IsTensorStridesColumnMajor(const std::shared_ptr<DataType>& type,
+ const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& strides) {
std::vector<int64_t> f_strides;
const auto& fw_type = checked_cast<const FixedWidthType&>(*type);
if (internal::ComputeColumnMajorStrides(fw_type, shape, &f_strides).ok()) {
@@ -137,9 +164,9 @@ inline bool IsTensorStridesColumnMajor(const
std::shared_ptr<DataType>& type,
}
}
-inline Status CheckTensorValidity(const std::shared_ptr<DataType>& type,
- const std::shared_ptr<Buffer>& data,
- const std::vector<int64_t>& shape) {
+Status CheckTensorValidity(const std::shared_ptr<DataType>& type,
+ const std::shared_ptr<Buffer>& data,
+ const std::vector<int64_t>& shape) {
if (!type) {
return Status::Invalid("Null type is supplied");
}
@@ -193,8 +220,8 @@ Status CheckTensorStridesValidity(const
std::shared_ptr<Buffer>& data,
}
return Status::OK();
}
-
} // namespace
+} // namespace internal
namespace internal {
@@ -532,11 +559,11 @@ bool Tensor::is_contiguous() const {
}
bool Tensor::is_row_major() const {
- return IsTensorStridesRowMajor(type_, shape_, strides_);
+ return internal::IsTensorStridesRowMajor(type_, shape_, strides_);
}
bool Tensor::is_column_major() const {
- return IsTensorStridesColumnMajor(type_, shape_, strides_);
+ return internal::IsTensorStridesColumnMajor(type_, shape_, strides_);
}
Type::type Tensor::type_id() const { return type_->id(); }
diff --git a/cpp/src/arrow/tensor.h b/cpp/src/arrow/tensor.h
index f3270313434..2917905049c 100644
--- a/cpp/src/arrow/tensor.h
+++ b/cpp/src/arrow/tensor.h
@@ -71,6 +71,15 @@ bool IsTensorStridesContiguous(const
std::shared_ptr<DataType>& type,
const std::vector<int64_t>& shape,
const std::vector<int64_t>& strides);
+/// Compute the size needed to store the tensor with the given strides and
shape.
+///
+/// If the strides are in number of element, pass `elem_size=1` to compute the
buffer size
+/// in the number of elements. If the strides are in bytes, pass the element
size in byte
+/// to `elem_size` and get the result in bytes.
+ARROW_EXPORT
+Result<int64_t> ComputeTensorSize(std::span<const int64_t> shape,
+ std::span<const int64_t> strides, int64_t
elem_size);
+
ARROW_EXPORT
Status ValidateTensorParameters(const std::shared_ptr<DataType>& type,
const std::shared_ptr<Buffer>& data,
diff --git a/docs/source/python/dlpack.rst b/docs/source/python/dlpack.rst
index 6e74cd5c82c..4e102ce8b0a 100644
--- a/docs/source/python/dlpack.rst
+++ b/docs/source/python/dlpack.rst
@@ -41,31 +41,51 @@ and more about DLPack in the
Implementation of DLPack in PyArrow
-----------------------------------
-The producing side of the DLPack Protocol is implemented for ``pa.Array``
-and can be used to interchange data between PyArrow and other tensor
-libraries. Supported data types are integer, unsigned integer and float. The
-protocol has no missing data support meaning PyArrow arrays with
-missing values cannot be transferred through the DLPack
-protocol. Currently, the Arrow implementation of the protocol only supports
+The protocol is implemented for ``pa.Array`` and ``pa.Tensor`` with different
behaviors.
+``pa.Tensor`` can produce and consume all shapes and strides of a generic
DLPack tensor.
+``pa.Array`` on the other hand is purposely limited to produce and consume
1-dimensional
+contiguous tensors (where the only dimension is the array's length).
+The only exception is ``pa.FixedShapeTensorArray``, which is designed to
represent
+tensors and supports more generic shapes and strides.
+It can produce and consume DLPack tensors whose outermost dimension has the
largest
+stride, that dimension being mapped to the array's length.
+
+For both ``pa.Tensor`` and ``pa.Array``, only numeric data types are
supported: integer,
+unsigned integer and float.
+
+Some array types can be understood as some form of tensor.
+For instance, a nested fixed size list of a numeric data type has the same
memory
+representation as a row major tensor.
+It is possible to get a (zero-copy) tensor from such an array using
+``array.to_tensor()``, and then use DLPack on the resulting tensor.
+
+The DLPack protocol fails on arrays with nulls, though these can be ignored
with
+an explicit conversion to a tensor using ``array.to_tensor(allow_nulls=True)``.
+In that case, the null entries hold an unspecified value.
+This is free, compared to ``pa.compute.fill_null`` which explicitly modifies
the
+array data to replace the null values.
+
+Currently, the Arrow implementation of the protocol only supports
data on a CPU device.
Data interchange syntax of the protocol includes
-1. ``from_dlpack(x)``: consuming an array object that implements a
- ``__dlpack__`` method and creating a new array while sharing the
+1. ``from_dlpack(x, /, *, device=None, copy=None)``: consuming an array object
that
+ implements a ``__dlpack__`` method and creating a new array while sharing
the
memory.
-2. ``__dlpack__(self, stream=None)`` and ``__dlpack_device__``:
+2. ``__dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None)``
+ and ``__dlpack_device__``:
producing a PyCapsule with the DLPack struct which is called from
within ``from_dlpack(x)``.
-
-PyArrow implements the second part of the protocol
-(``__dlpack__(self, stream=None)`` and ``__dlpack_device__``) and can
-thus be consumed by libraries implementing ``from_dlpack``.
+ This method is intended for library authors.
Examples
--------
+Producing
+~~~~~~~~~
+
Convert a PyArrow CPU array into a NumPy array:
.. code-block:: python
@@ -101,3 +121,81 @@ Convert a PyArrow CPU array into a JAX array:
Array([2, 0, 2, 4], dtype=int32)
>>> jax.dlpack.from_dlpack(array) # doctest: +SKIP
Array([2, 0, 2, 4], dtype=int32)
+
+Arrays with a tensor memory layout, such as fixed size lists of a numeric
type, need an
+explicit conversion to a ``pa.Tensor``, which exports its full
multi-dimensional shape:
+
+.. code-block:: python
+
+ >>> list_array = pa.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
+ ... pa.list_(pa.float64(), 2))
+ >>> list_array.to_tensor()
+ <pyarrow.Tensor>
+ type: double
+ shape: (3, 2)
+ strides: (16, 8)
+ >>> np.from_dlpack(list_array.to_tensor())
+ array([[1., 2.],
+ [3., 4.],
+ [5., 6.]])
+
+A ``pa.FixedShapeTensorArray`` exports directly, the array length becoming the
outermost
+dimension, followed by the shape of the element tensors:
+
+.. code-block:: python
+
+ >>> nested = pa.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+ ... pa.list_(pa.list_(pa.int32(), 2), 2))
+ >>> tensor_array = pa.FixedShapeTensorArray.from_tensor(nested.to_tensor())
+ >>> tensor_array.type
+ FixedShapeTensorType(extension<arrow.fixed_shape_tensor[value_type=int32,
shape=[2,2], permutation=[0,1]]>)
+ >>> np.from_dlpack(tensor_array).shape
+ (2, 2, 2)
+
+Arrays with nulls are rejected, unless the conversion to a tensor explicitly
allows
+them, in which case the null entries hold an unspecified value:
+
+.. code-block:: python
+
+ >>> array_with_nulls = pa.array([2, None, 4], pa.int32())
+ >>> np.from_dlpack(array_with_nulls)
+ Traceback (most recent call last):
+ ...
+ pyarrow.lib.ArrowTypeError: Can only use DLPack on arrays with no nulls.
+ >>> np.from_dlpack(array_with_nulls.to_tensor(allow_nulls=True))
+ array([2, ..., 4], dtype=int32)
+
+Consuming
+~~~~~~~~~
+
+Any object implementing the DLPack protocol can be imported, without copying
the data:
+
+.. code-block:: python
+
+ >>> pa.Array.from_dlpack(np.array([2, 0, 2, 4]))
+ <pyarrow.lib.Int64Array object at ...>
+ [
+ 2,
+ 0,
+ 2,
+ 4
+ ]
+ >>> pa.Tensor.from_dlpack(np.array([[2, 0], [2, 4]], np.int32))
+ <pyarrow.Tensor>
+ type: int32
+ shape: (2, 2)
+ strides: (8, 4)
+
+``pa.Array.from_dlpack`` only accepts 1-dimensional contiguous tensors.
+Multi-dimensional data can be imported as a ``pa.FixedShapeTensorArray``, the
outermost
+dimension becoming the length of the array:
+
+.. code-block:: python
+
+ >>> array = pa.FixedShapeTensorArray.from_dlpack(
+ ... np.arange(12, dtype=np.int32).reshape(3, 2, 2)
+ ... )
+ >>> array.type
+ FixedShapeTensorType(extension<arrow.fixed_shape_tensor[value_type=int32,
shape=[2,2], permutation=[0,1]]>)
+ >>> len(array)
+ 3
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index 691623b87f0..3060c533255 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -15,7 +15,12 @@
# specific language governing permissions and limitations
# under the License.
-from cpython.pycapsule cimport PyCapsule_CheckExact, PyCapsule_GetPointer,
PyCapsule_New
+from cpython.pycapsule cimport (
+ PyCapsule_CheckExact,
+ PyCapsule_GetPointer,
+ PyCapsule_New,
+ PyCapsule_SetName,
+)
from collections.abc import Sequence
import os
@@ -2269,6 +2274,54 @@ 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.
+ Only 1-dimensional contiguous tensors are accepted as input.
+ For multi-dimensional tensors, use `Tensor.from_dlpack` or
+ `FixedShapeTensorArray.from_dlpack`.
+
+ 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)
+ carray = GetResultValue(result)
+ return pyarrow_wrap_array(carray)
+
def __dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None):
"""
Export a primitive array as a DLPack capsule.
@@ -4969,6 +5022,32 @@ cdef class FixedShapeTensorArray(ExtensionArray):
return self.to_tensor().to_numpy()
+ @staticmethod
+ def from_tensor(Tensor tensor not None):
+ """
+ Convert a pyarrow.Tensor to a fixed shape tensor extension array.
+
+ The first dimension of the tensor becomes the length of the fixed shape
+ tensor array and the remaining dimensions the shape of the individual
+ tensors. If the tensor provides strides, they are used to determine the
+ dimension permutation, otherwise row-major layout is assumed.
+
+ Parameters
+ ----------
+ tensor : pyarrow.Tensor
+
+ Returns
+ -------
+ FixedShapeTensorArray
+ """
+ cdef shared_ptr[CFixedShapeTensorArray] c_array
+
+ with nogil:
+ c_array = GetResultValue(
+ CFixedShapeTensorArray.FromTensor(tensor.sp_tensor))
+
+ return pyarrow_wrap_array(<shared_ptr[CArray]> c_array)
+
@staticmethod
def from_numpy_ndarray(obj, dim_names=None):
"""
@@ -5045,6 +5124,41 @@ cdef class FixedShapeTensorArray(ExtensionArray):
FixedSizeListArray.from_arrays(values, shape[1:].prod())
)
+ @staticmethod
+ def from_dlpack(x, /, *, device=None, copy=None):
+ """
+ Construct a FixedShapeTensorArray from an object implementing the
DLPack
+ protocol.
+
+ The outermost dimension of the input becomes the length of the tensor
+ array, and the remaining dimensions the shape of the individual
tensors.
+ The outermost dimension must have the largest stride.
+
+ 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
+ -------
+ FixedShapeTensorArray
+ An array housing the data from the input object, potentially
+ as a copy or view.
+
+ """
+ return FixedShapeTensorArray.from_tensor(
+ Tensor.from_dlpack(x, device=device, copy=copy))
+
def __dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None):
"""
Export a tensor array as a DLPack capsule.
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index d592d4024aa..16784b51ecc 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -1460,6 +1460,10 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
cdef extern from "arrow/c/dlpack_abi.h" nogil:
+ ctypedef struct DLPackVersion:
+ uint32_t major
+ uint32_t minor
+
ctypedef enum DLDeviceType:
kDLCPU = 1
@@ -1475,6 +1479,8 @@ cdef extern from "arrow/c/dlpack_abi.h" nogil:
cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil:
+ const DLPackVersion DLPACK_VERSION" arrow::dlpack::kVersion"
+
CResult[DLManagedTensor*] ExportArrayToDLPack" arrow::dlpack::ExportArray"(
const shared_ptr[CArray]& arr)
CResult[DLManagedTensor*] ExportTensorToDLPack"
arrow::dlpack::ExportTensor"(
@@ -1490,6 +1496,13 @@ cdef extern from "arrow/c/dlpack.h" namespace
"arrow::dlpack" nogil:
CResult[DLDevice] ExportDevice(const shared_ptr[CArray]& arr)
CResult[DLDevice] ExportDevice(const shared_ptr[CTensor]& tensor)
+ CResult[shared_ptr[CArray]] \
+ ImportArrayVersionedFromDLPack" arrow::dlpack::ImportArrayVersioned"(
+ DLManagedTensorVersioned* raw)
+ CResult[shared_ptr[CTensor]] \
+ ImportTensorVersionedFromDLPack" arrow::dlpack::ImportTensorVersioned"(
+ DLManagedTensorVersioned* raw)
+
cdef extern from "arrow/builder.h" namespace "arrow" nogil:
@@ -3092,6 +3105,13 @@ cdef extern from "arrow/extension/fixed_shape_tensor.h"
namespace "arrow::extens
const vector[int64_t] permutation()
const vector[c_string] dim_names()
+ cdef cppclass CFixedShapeTensorArray \
+ " arrow::extension::FixedShapeTensorArray"(CExtensionArray):
+
+ @staticmethod
+ CResult[shared_ptr[CFixedShapeTensorArray]] FromTensor(
+ const shared_ptr[CTensor]& tensor)
+
cdef extern from "arrow/extension/opaque.h" namespace "arrow::extension" nogil:
cdef cppclass COpaqueType \
diff --git a/python/pyarrow/tensor.pxi b/python/pyarrow/tensor.pxi
index 521ee0c3f44..7f006f29417 100644
--- a/python/pyarrow/tensor.pxi
+++ b/python/pyarrow/tensor.pxi
@@ -18,6 +18,12 @@
# Avoid name clash with `pa.struct` function
import struct as _struct
+from cpython.pycapsule cimport (
+ PyCapsule_CheckExact,
+ PyCapsule_GetPointer,
+ PyCapsule_SetName,
+)
+
cdef class Tensor(_Weakrefable):
"""
@@ -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)
+ ctensor = GetResultValue(result)
+ return pyarrow_wrap_tensor(ctensor)
+
+ def __dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None):
"""
Export a Tensor as a DLPack capsule.
diff --git a/python/pyarrow/tests/test_dlpack.py
b/python/pyarrow/tests/test_dlpack.py
index e3cc2fd3e9e..0c3c082fd6e 100644
--- a/python/pyarrow/tests/test_dlpack.py
+++ b/python/pyarrow/tests/test_dlpack.py
@@ -374,3 +374,133 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@requires_numpy_version("2.1.0")
+@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):
+ def make_array():
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ array = base[::2, 1::2]
+ assert not array.flags['C_CONTIGUOUS']
+ return array
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ tensor = pa.Tensor.from_dlpack(make_array())
+ assert isinstance(tensor, pa.Tensor)
+ gc.collect() # Attempts to free input array memory
+ np.testing.assert_array_equal(tensor.to_numpy(), make_array(), strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@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):
+ expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
+ arr = pa.Array.from_dlpack(expected)
+ arr.validate(full=True)
+ assert isinstance(arr, pa.Array)
+ np.testing.assert_array_equal(arr.to_numpy(), expected, strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@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_fixed_shape_tensor_array_from_dlpack(np_type):
+ source = np.arange(12, dtype=np_type).reshape((3, 2, 2))
+ arr = pa.FixedShapeTensorArray.from_dlpack(source)
+ arr.validate(full=True)
+ assert arr.type == pa.fixed_shape_tensor(pa.from_numpy_dtype(np_type), [2,
2])
+ assert arr.to_pylist() == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
+
+ # Zero-copy import: mutating the source is visible through the array.
+ source[0, 0, 0] = 100
+ assert arr.to_pylist()[0] == [100, 1, 2, 3]
+
+ copied = pa.FixedShapeTensorArray.from_dlpack(source, copy=True)
+ source[0, 0, 0] = 0
+ assert copied.to_pylist()[0] == [100, 1, 2, 3]
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type', [np.uint8, np.int32, np.float64])
+def test_fixed_shape_tensor_array_from_dlpack_transposed(np_type):
+ source = np.arange(12, dtype=np_type).reshape((3, 2, 2)).transpose(0, 2, 1)
+ arr = pa.FixedShapeTensorArray.from_dlpack(source)
+ arr.validate(full=True)
+ assert arr.type == pa.fixed_shape_tensor(
+ pa.from_numpy_dtype(np_type), [2, 2], permutation=[1, 0])
+ np.testing.assert_array_equal(arr.to_numpy_ndarray(), source)
+
+ # Zero-copy import: mutating the source is visible through the array.
+ source[0, 0, 0] = 100
+ np.testing.assert_array_equal(arr.to_numpy_ndarray(), source)
+
+ copied = pa.FixedShapeTensorArray.from_dlpack(source, copy=True)
+ expected = source.copy()
+ source[0, 0, 0] = 0
+ np.testing.assert_array_equal(copied.to_numpy_ndarray(), expected)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
+def test_fixed_shape_tensor_array_from_dlpack_not_first_major():
+ # The outermost dimension indexes the tensor elements, so it must remain
+ # the major one.
+ source = np.arange(12, dtype=np.int32).reshape((3, 2, 2)).transpose(1, 0,
2)
+ with pytest.raises(pa.ArrowInvalid,
+ match="Only first-major tensors can be zero-copy"):
+ pa.FixedShapeTensorArray.from_dlpack(source)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
+def test_from_dlpack_zero_copy():
+ expected = np.array([1, 2, 3], dtype=np.int64)
+ tensor = pa.Tensor.from_dlpack(expected)
+ result = tensor.to_numpy()
+ # Zero-copy import: mutating the source is visible through the tensor.
+ expected[0] = 100
+ assert result[0] == 100
+ # Same for mutating the result
+ result[1] = 42
+ assert expected[1] == 42
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
+def test_from_dlpack_explicit_copy():
+ expected = np.array([1, 2, 3], dtype=np.int64)
+ tensor = pa.Tensor.from_dlpack(expected, copy=True)
+ result = tensor.to_numpy()
+ expected[0] = 100
+ # The data was copied, so mutating the source is not visible.
+ assert result[0] == 1
+
+
+def test_from_dlpack_no_dlpack_method():
+ with pytest.raises(AttributeError):
+ pa.Tensor.from_dlpack(object())
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
+def test_array_from_dlpack_multi_dim_not_supported():
+ expected = np.arange(6, dtype=np.int32).reshape((2, 3))
+ with pytest.raises(
+ pa.ArrowInvalid,
+ match="Only contiguous one dimensional tensor can be imported as
arrays",
+ ):
+ pa.Array.from_dlpack(expected)
diff --git a/python/pyarrow/tests/test_extension_type.py
b/python/pyarrow/tests/test_extension_type.py
index bdd898767b0..a9d15cfec10 100644
--- a/python/pyarrow/tests/test_extension_type.py
+++ b/python/pyarrow/tests/test_extension_type.py
@@ -1708,6 +1708,37 @@ def test_tensor_class_methods(np_type_str):
assert result.to_tensor().strides == (12 * bw, 1 * bw, 3 * bw, 6 * bw)
[email protected]
[email protected](
+ ("transpose", "permutation"),
+ [(False, [0, 1]), (True, [1, 0])]
+)
+def test_tensor_array_from_tensor(transpose, permutation):
+ arr = np.arange(24, dtype=np.int32).reshape(2, 3, 4)
+ arr = arr.transpose(0, 2, 1) if transpose else arr
+
+ result = pa.FixedShapeTensorArray.from_tensor(pa.Tensor.from_numpy(arr))
+ result.validate(full=True)
+
+ assert isinstance(result.type, pa.FixedShapeTensorType)
+ assert result.type.value_type == pa.int32()
+ # Shape is in physical order (unpermuted)
+ assert result.type.shape == [3, 4]
+ assert result.type.permutation == permutation
+ assert len(result) == 2
+ np.testing.assert_array_equal(result.to_numpy_ndarray(), arr)
+
+
[email protected]
[email protected]("permutation", [(1, 0, 2), (1, 2, 0), (2, 1, 0)])
+def test_tensor_array_from_tensor_not_first_major(permutation):
+ arr = np.arange(24, dtype=np.int32).reshape(2, 3,
4).transpose(*permutation)
+
+ with pytest.raises(pa.ArrowInvalid,
+ match="Only first-major tensors can be zero-copy"):
+ pa.FixedShapeTensorArray.from_tensor(pa.Tensor.from_numpy(arr))
+
+
@pytest.mark.numpy
@pytest.mark.parametrize("np_type_str", ("int8", "int64", "float32"))
def test_tensor_array_from_numpy(np_type_str):
