This is an automated email from the ASF dual-hosted git repository.
apitrou 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 6c326db6a5 GH-33984: [C++][Python] DLPack implementation for Arrow
Arrays (producer) (#38472)
6c326db6a5 is described below
commit 6c326db6a5686a78bc77be662b61236ddbfc66dc
Author: Alenka Frim <[email protected]>
AuthorDate: Tue Dec 19 19:58:29 2023 +0100
GH-33984: [C++][Python] DLPack implementation for Arrow Arrays (producer)
(#38472)
### Rationale for this change
DLPack is selected for Array API protocol so it is important to have it
implemented for Arrow/PyArrow Arrays also. This is possible for primitive type
arrays (int, uint and float) with no validity buffer. Device support is not in
scope of this PR (CPU only).
### What changes are included in this PR?
- `ExportArray` and `ExportDevice` methods on Arrow C++ Arrays
- `__dlpack__` method on the base PyArrow Array class exposing
`ExportArray` method
- `__dlpack_device__` method on the base PyArrow Array class exposing
`ExportDevice` method
### Are these changes tested?
Yes, tests are added to `dlpack_test.cc` and `test_array.py`.
### Are there any user-facing changes?
No.
* Closes: #33984
Lead-authored-by: AlenkaF <[email protected]>
Co-authored-by: Alenka Frim <[email protected]>
Co-authored-by: Antoine Pitrou <[email protected]>
Co-authored-by: Joris Van den Bossche <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/CMakeLists.txt | 1 +
cpp/src/arrow/c/CMakeLists.txt | 1 +
cpp/src/arrow/c/dlpack.cc | 133 ++++++++++++
cpp/src/arrow/c/dlpack.h | 51 +++++
cpp/src/arrow/c/dlpack_abi.h | 321 ++++++++++++++++++++++++++++
cpp/src/arrow/c/dlpack_test.cc | 129 +++++++++++
dev/release/rat_exclude_files.txt | 1 +
docs/source/python/dlpack.rst | 93 ++++++++
docs/source/python/index.rst | 1 +
docs/source/python/interchange_protocol.rst | 6 +-
python/pyarrow/_dlpack.pxi | 46 ++++
python/pyarrow/array.pxi | 38 ++++
python/pyarrow/includes/libarrow.pxd | 19 ++
python/pyarrow/lib.pyx | 3 +
python/pyarrow/tests/test_dlpack.py | 142 ++++++++++++
15 files changed, 982 insertions(+), 3 deletions(-)
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
index 46a7aa9106..00947c6275 100644
--- a/cpp/src/arrow/CMakeLists.txt
+++ b/cpp/src/arrow/CMakeLists.txt
@@ -192,6 +192,7 @@ set(ARROW_SRCS
type_traits.cc
visitor.cc
c/bridge.cc
+ c/dlpack.cc
io/buffered.cc
io/caching.cc
io/compressed.cc
diff --git a/cpp/src/arrow/c/CMakeLists.txt b/cpp/src/arrow/c/CMakeLists.txt
index 3765477ba0..81a81cd3f1 100644
--- a/cpp/src/arrow/c/CMakeLists.txt
+++ b/cpp/src/arrow/c/CMakeLists.txt
@@ -16,6 +16,7 @@
# under the License.
add_arrow_test(bridge_test PREFIX "arrow-c")
+add_arrow_test(dlpack_test)
add_arrow_benchmark(bridge_benchmark)
diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc
new file mode 100644
index 0000000000..13ee2761b0
--- /dev/null
+++ b/cpp/src/arrow/c/dlpack.cc
@@ -0,0 +1,133 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/c/dlpack.h"
+
+#include "arrow/array/array_base.h"
+#include "arrow/c/dlpack_abi.h"
+#include "arrow/device.h"
+#include "arrow/type.h"
+#include "arrow/type_traits.h"
+
+namespace arrow::dlpack {
+
+namespace {
+
+Result<DLDataType> GetDLDataType(const DataType& type) {
+ DLDataType dtype;
+ dtype.lanes = 1;
+ dtype.bits = type.bit_width();
+ switch (type.id()) {
+ case Type::INT8:
+ case Type::INT16:
+ case Type::INT32:
+ case Type::INT64:
+ dtype.code = DLDataTypeCode::kDLInt;
+ return dtype;
+ case Type::UINT8:
+ case Type::UINT16:
+ case Type::UINT32:
+ case Type::UINT64:
+ dtype.code = DLDataTypeCode::kDLUInt;
+ return dtype;
+ case Type::HALF_FLOAT:
+ case Type::FLOAT:
+ case Type::DOUBLE:
+ dtype.code = DLDataTypeCode::kDLFloat;
+ return dtype;
+ case Type::BOOL:
+ // DLPack supports byte-packed boolean values
+ return Status::TypeError("Bit-packed boolean data type not supported by
DLPack.");
+ default:
+ return Status::TypeError("DataType is not compatible with DLPack spec: ",
+ type.ToString());
+ }
+}
+
+struct ManagerCtx {
+ std::shared_ptr<ArrayData> array;
+ DLManagedTensor tensor;
+};
+
+} // namespace
+
+Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr) {
+ // Define DLDevice struct nad check if array type is supported
+ // by the DLPack protocol at the same time. Raise TypeError if not.
+ // Supported data types: int, uint, float with no validity buffer.
+ ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(arr))
+
+ // Define the DLDataType struct
+ const DataType& type = *arr->type();
+ std::shared_ptr<ArrayData> data = arr->data();
+ ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type));
+
+ // Create ManagerCtx that will serve as the owner of the DLManagedTensor
+ std::unique_ptr<ManagerCtx> ctx(new ManagerCtx);
+
+ // Define the data pointer to the DLTensor
+ // If array is of length 0, data pointer should be NULL
+ if (arr->length() == 0) {
+ ctx->tensor.dl_tensor.data = NULL;
+ } else {
+ const auto data_offset = data->offset * type.byte_width();
+ ctx->tensor.dl_tensor.data =
+ const_cast<uint8_t*>(data->buffers[1]->data() + data_offset);
+ }
+
+ ctx->tensor.dl_tensor.device = device;
+ ctx->tensor.dl_tensor.ndim = 1;
+ ctx->tensor.dl_tensor.dtype = dlpack_type;
+ ctx->tensor.dl_tensor.shape = const_cast<int64_t*>(&data->length);
+ ctx->tensor.dl_tensor.strides = NULL;
+ ctx->tensor.dl_tensor.byte_offset = 0;
+
+ ctx->array = std::move(data);
+ ctx->tensor.manager_ctx = ctx.get();
+ ctx->tensor.deleter = [](struct DLManagedTensor* self) {
+ delete reinterpret_cast<ManagerCtx*>(self->manager_ctx);
+ };
+ return &ctx.release()->tensor;
+}
+
+Result<DLDevice> ExportDevice(const std::shared_ptr<Array>& arr) {
+ // Check if array is supported by the DLPack protocol.
+ if (arr->null_count() > 0) {
+ return Status::TypeError("Can only use DLPack on arrays with no nulls.");
+ }
+ const DataType& type = *arr->type();
+ if (type.id() == Type::BOOL) {
+ return Status::TypeError("Bit-packed boolean data type not supported by
DLPack.");
+ }
+ if (!is_integer(type.id()) && !is_floating(type.id())) {
+ return Status::TypeError("DataType is not compatible with DLPack spec: ",
+ type.ToString());
+ }
+
+ // Define DLDevice struct
+ DLDevice device;
+ if (arr->data()->buffers[1]->device_type() == DeviceAllocationType::kCPU) {
+ device.device_id = 0;
+ device.device_type = DLDeviceType::kDLCPU;
+ return device;
+ } else {
+ return Status::NotImplemented(
+ "DLPack support is implemented only for buffers on CPU device.");
+ }
+}
+
+} // namespace arrow::dlpack
diff --git a/cpp/src/arrow/c/dlpack.h b/cpp/src/arrow/c/dlpack.h
new file mode 100644
index 0000000000..d11ccfc1fd
--- /dev/null
+++ b/cpp/src/arrow/c/dlpack.h
@@ -0,0 +1,51 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include "arrow/array/array_base.h"
+#include "arrow/c/dlpack_abi.h"
+
+namespace arrow::dlpack {
+
+/// \brief Export Arrow array as DLPack tensor.
+///
+/// DLMangedTensor is produced as defined by the DLPack protocol,
+/// see https://dmlc.github.io/dlpack/latest/.
+///
+/// Data types for which the protocol is supported are
+/// integer and floating-point data types.
+///
+/// DLPack protocol only supports arrays with one contiguous
+/// memory region which means Arrow Arrays with validity buffers
+/// are not supported.
+///
+/// \param[in] arr Arrow array
+/// \return DLManagedTensor struct
+ARROW_EXPORT
+Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr);
+
+/// \brief Get DLDevice with enumerator specifying the
+/// type of the device data is stored on and index of the
+/// device which is 0 by default for CPU.
+///
+/// \param[in] arr Arrow array
+/// \return DLDevice struct
+ARROW_EXPORT
+Result<DLDevice> ExportDevice(const std::shared_ptr<Array>& arr);
+
+} // namespace arrow::dlpack
diff --git a/cpp/src/arrow/c/dlpack_abi.h b/cpp/src/arrow/c/dlpack_abi.h
new file mode 100644
index 0000000000..4af557a7ed
--- /dev/null
+++ b/cpp/src/arrow/c/dlpack_abi.h
@@ -0,0 +1,321 @@
+// Taken from:
+//
https://github.com/dmlc/dlpack/blob/ca4d00ad3e2e0f410eeab3264d21b8a39397f362/include/dlpack/dlpack.h
+/*!
+ * Copyright (c) 2017 by Contributors
+ * \file dlpack.h
+ * \brief The common header of DLPack.
+ */
+#ifndef DLPACK_DLPACK_H_
+#define DLPACK_DLPACK_H_
+
+/**
+ * \brief Compatibility with C++
+ */
+#ifdef __cplusplus
+#define DLPACK_EXTERN_C extern "C"
+#else
+#define DLPACK_EXTERN_C
+#endif
+
+/*! \brief The current major version of dlpack */
+#define DLPACK_MAJOR_VERSION 1
+
+/*! \brief The current minor version of dlpack */
+#define DLPACK_MINOR_VERSION 0
+
+/*! \brief DLPACK_DLL prefix for windows */
+#ifdef _WIN32
+#ifdef DLPACK_EXPORTS
+#define DLPACK_DLL __declspec(dllexport)
+#else
+#define DLPACK_DLL __declspec(dllimport)
+#endif
+#else
+#define DLPACK_DLL
+#endif
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*!
+ * \brief The DLPack version.
+ *
+ * A change in major version indicates that we have changed the
+ * data layout of the ABI - DLManagedTensorVersioned.
+ *
+ * A change in minor version indicates that we have added new
+ * code, such as a new device type, but the ABI is kept the same.
+ *
+ * If an obtained DLPack tensor has a major version that disagrees
+ * with the version number specified in this header file
+ * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter
+ * (and it is safe to do so). It is not safe to access any other fields
+ * as the memory layout will have changed.
+ *
+ * In the case of a minor version mismatch, the tensor can be safely used as
+ * long as the consumer knows how to interpret all fields. Minor version
+ * updates indicate the addition of enumeration values.
+ */
+typedef struct {
+ /*! \brief DLPack major version. */
+ uint32_t major;
+ /*! \brief DLPack minor version. */
+ uint32_t minor;
+} DLPackVersion;
+
+/*!
+ * \brief The device type in DLDevice.
+ */
+#ifdef __cplusplus
+typedef enum : int32_t {
+#else
+typedef enum {
+#endif
+ /*! \brief CPU device */
+ kDLCPU = 1,
+ /*! \brief CUDA GPU device */
+ kDLCUDA = 2,
+ /*!
+ * \brief Pinned CUDA CPU memory by cudaMallocHost
+ */
+ kDLCUDAHost = 3,
+ /*! \brief OpenCL devices. */
+ kDLOpenCL = 4,
+ /*! \brief Vulkan buffer for next generation graphics. */
+ kDLVulkan = 7,
+ /*! \brief Metal for Apple GPU. */
+ kDLMetal = 8,
+ /*! \brief Verilog simulator buffer */
+ kDLVPI = 9,
+ /*! \brief ROCm GPUs for AMD GPUs */
+ kDLROCM = 10,
+ /*!
+ * \brief Pinned ROCm CPU memory allocated by hipMallocHost
+ */
+ kDLROCMHost = 11,
+ /*!
+ * \brief Reserved extension device type,
+ * used for quickly test extension device
+ * The semantics can differ depending on the implementation.
+ */
+ kDLExtDev = 12,
+ /*!
+ * \brief CUDA managed/unified memory allocated by cudaMallocManaged
+ */
+ kDLCUDAManaged = 13,
+ /*!
+ * \brief Unified shared memory allocated on a oneAPI non-partititioned
+ * device. Call to oneAPI runtime is required to determine the device
+ * type, the USM allocation type and the sycl context it is bound to.
+ *
+ */
+ kDLOneAPI = 14,
+ /*! \brief GPU support for next generation WebGPU standard. */
+ kDLWebGPU = 15,
+ /*! \brief Qualcomm Hexagon DSP */
+ kDLHexagon = 16,
+} DLDeviceType;
+
+/*!
+ * \brief A Device for Tensor and operator.
+ */
+typedef struct {
+ /*! \brief The device type used in the device. */
+ DLDeviceType device_type;
+ /*!
+ * \brief The device index.
+ * For vanilla CPU memory, pinned memory, or managed memory, this is set to
0.
+ */
+ int32_t device_id;
+} DLDevice;
+
+/*!
+ * \brief The type code options DLDataType.
+ */
+typedef enum {
+ /*! \brief signed integer */
+ kDLInt = 0U,
+ /*! \brief unsigned integer */
+ kDLUInt = 1U,
+ /*! \brief IEEE floating point */
+ kDLFloat = 2U,
+ /*!
+ * \brief Opaque handle type, reserved for testing purposes.
+ * Frameworks need to agree on the handle data type for the exchange to be
well-defined.
+ */
+ kDLOpaqueHandle = 3U,
+ /*! \brief bfloat16 */
+ kDLBfloat = 4U,
+ /*!
+ * \brief complex number
+ * (C/C++/Python layout: compact struct per complex number)
+ */
+ kDLComplex = 5U,
+ /*! \brief boolean */
+ kDLBool = 6U,
+} DLDataTypeCode;
+
+/*!
+ * \brief The data type the tensor can hold. The data type is assumed to
follow the
+ * native endian-ness. An explicit error message should be raised when
attempting to
+ * export an array with non-native endianness
+ *
+ * Examples
+ * - float: type_code = 2, bits = 32, lanes = 1
+ * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4
+ * - int8: type_code = 0, bits = 8, lanes = 1
+ * - std::complex<float>: type_code = 5, bits = 64, lanes = 1
+ * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library
convention,
+ * the underlying storage size of bool is 8 bits)
+ */
+typedef struct {
+ /*!
+ * \brief Type code of base types.
+ * We keep it uint8_t instead of DLDataTypeCode for minimal memory
+ * footprint, but the value should be one of DLDataTypeCode enum values.
+ * */
+ uint8_t code;
+ /*!
+ * \brief Number of bits, common choices are 8, 16, 32.
+ */
+ uint8_t bits;
+ /*! \brief Number of lanes in the type, used for vector types. */
+ uint16_t lanes;
+} DLDataType;
+
+/*!
+ * \brief Plain C Tensor object, does not manage memory.
+ */
+typedef struct {
+ /*!
+ * \brief The data pointer points to the allocated data. This will be CUDA
+ * device pointer or cl_mem handle in OpenCL. It may be opaque on some device
+ * types. This pointer is always aligned to 256 bytes as in CUDA. The
+ * `byte_offset` field should be used to point to the beginning of the data.
+ *
+ * Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,
+ * TVM, perhaps others) do not adhere to this 256 byte aligment requirement
+ * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed
+ * (after which this note will be updated); at the moment it is recommended
+ * to not rely on the data pointer being correctly aligned.
+ *
+ * For given DLTensor, the size of memory required to store the contents of
+ * data is calculated as follows:
+ *
+ * \code{.c}
+ * static inline size_t GetDataSize(const DLTensor* t) {
+ * size_t size = 1;
+ * for (tvm_index_t i = 0; i < t->ndim; ++i) {
+ * size *= t->shape[i];
+ * }
+ * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8;
+ * return size;
+ * }
+ * \endcode
+ */
+ void* data;
+ /*! \brief The device of the tensor */
+ DLDevice device;
+ /*! \brief Number of dimensions */
+ int32_t ndim;
+ /*! \brief The data type of the pointer*/
+ DLDataType dtype;
+ /*! \brief The shape of the tensor */
+ int64_t* shape;
+ /*!
+ * \brief strides of the tensor (in number of elements, not bytes)
+ * can be NULL, indicating tensor is compact and row-majored.
+ */
+ int64_t* strides;
+ /*! \brief The offset in bytes to the beginning pointer to data */
+ uint64_t byte_offset;
+} DLTensor;
+
+/*!
+ * \brief C Tensor object, manage memory of DLTensor. This data structure is
+ * intended to facilitate the borrowing of DLTensor by another framework. It
is
+ * not meant to transfer the tensor. When the borrowing framework doesn't need
+ * the tensor, it should call the deleter to notify the host that the resource
+ * is no longer needed.
+ *
+ * \note This data structure is used as Legacy DLManagedTensor
+ * in DLPack exchange and is deprecated after DLPack v0.8
+ * Use DLManagedTensorVersioned instead.
+ * This data structure may get renamed or deleted in future versions.
+ *
+ * \sa DLManagedTensorVersioned
+ */
+typedef struct DLManagedTensor {
+ /*! \brief DLTensor which is being memory managed */
+ DLTensor dl_tensor;
+ /*! \brief the context of the original host framework of DLManagedTensor in
+ * which DLManagedTensor is used in the framework. It can also be NULL.
+ */
+ void* manager_ctx;
+ /*!
+ * \brief Destructor - this should be called
+ * to destruct the manager_ctx which backs the DLManagedTensor. It can be
+ * NULL if there is no way for the caller to provide a reasonable destructor.
+ * The destructors deletes the argument self as well.
+ */
+ void (*deleter)(struct DLManagedTensor* self);
+} DLManagedTensor;
+
+// bit masks used in in the DLManagedTensorVersioned
+
+/*! \brief bit mask to indicate that the tensor is read only. */
+#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL)
+
+/*!
+ * \brief A versioned and managed C Tensor object, manage memory of DLTensor.
+ *
+ * This data structure is intended to facilitate the borrowing of DLTensor by
+ * another framework. It is not meant to transfer the tensor. When the
borrowing
+ * framework doesn't need the tensor, it should call the deleter to notify the
+ * host that the resource is no longer needed.
+ *
+ * \note This is the current standard DLPack exchange data structure.
+ */
+struct DLManagedTensorVersioned {
+ /*!
+ * \brief The API and ABI version of the current managed Tensor
+ */
+ DLPackVersion version;
+ /*!
+ * \brief the context of the original host framework.
+ *
+ * Stores DLManagedTensorVersioned is used in the
+ * framework. It can also be NULL.
+ */
+ void* manager_ctx;
+ /*!
+ * \brief Destructor.
+ *
+ * This should be called to destruct manager_ctx which holds the
+ * DLManagedTensorVersioned. It can be NULL if there is no way for the
caller to provide
+ * a reasonable destructor. The destructors deletes the argument self as
well.
+ */
+ void (*deleter)(struct DLManagedTensorVersioned* self);
+ /*!
+ * \brief Additional bitmask flags information about the tensor.
+ *
+ * By default the flags should be set to 0.
+ *
+ * \note Future ABI changes should keep everything until this field
+ * stable, to ensure that deleter can be correctly called.
+ *
+ * \sa DLPACK_FLAG_BITMASK_READ_ONLY
+ */
+ uint64_t flags;
+ /*! \brief DLTensor which is being memory managed */
+ DLTensor dl_tensor;
+};
+
+#ifdef __cplusplus
+} // DLPACK_EXTERN_C
+#endif
+#endif // DLPACK_DLPACK_H_
diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc
new file mode 100644
index 0000000000..3136506bf3
--- /dev/null
+++ b/cpp/src/arrow/c/dlpack_test.cc
@@ -0,0 +1,129 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+
+#include "arrow/array/array_base.h"
+#include "arrow/c/dlpack.h"
+#include "arrow/c/dlpack_abi.h"
+#include "arrow/memory_pool.h"
+#include "arrow/testing/gtest_util.h"
+
+namespace arrow::dlpack {
+
+class TestExportArray : public ::testing::Test {
+ public:
+ void SetUp() {}
+};
+
+void CheckDLTensor(const std::shared_ptr<Array>& arr,
+ const std::shared_ptr<DataType>& arrow_type,
+ DLDataTypeCode dlpack_type, int64_t length) {
+ ASSERT_OK_AND_ASSIGN(auto dlmtensor, arrow::dlpack::ExportArray(arr));
+ auto dltensor = dlmtensor->dl_tensor;
+
+ const auto byte_width = arr->type()->byte_width();
+ const auto start = arr->offset() * byte_width;
+ ASSERT_OK_AND_ASSIGN(auto sliced_buffer,
+ SliceBufferSafe(arr->data()->buffers[1], start));
+ ASSERT_EQ(sliced_buffer->data(), dltensor.data);
+
+ ASSERT_EQ(0, dltensor.byte_offset);
+ ASSERT_EQ(NULL, dltensor.strides);
+ ASSERT_EQ(length, dltensor.shape[0]);
+ ASSERT_EQ(1, dltensor.ndim);
+
+ ASSERT_EQ(dlpack_type, dltensor.dtype.code);
+
+ ASSERT_EQ(arrow_type->bit_width(), dltensor.dtype.bits);
+ ASSERT_EQ(1, dltensor.dtype.lanes);
+ ASSERT_EQ(DLDeviceType::kDLCPU, dltensor.device.device_type);
+ ASSERT_EQ(0, dltensor.device.device_id);
+
+ ASSERT_OK_AND_ASSIGN(auto device, arrow::dlpack::ExportDevice(arr));
+ ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type);
+ ASSERT_EQ(0, device.device_id);
+
+ dlmtensor->deleter(dlmtensor);
+}
+
+TEST_F(TestExportArray, TestSupportedArray) {
+ const std::vector<std::pair<std::shared_ptr<DataType>, DLDataTypeCode>>
cases = {
+ {int8(), DLDataTypeCode::kDLInt},
+ {uint8(), DLDataTypeCode::kDLUInt},
+ {
+ int16(),
+ DLDataTypeCode::kDLInt,
+ },
+ {uint16(), DLDataTypeCode::kDLUInt},
+ {
+ int32(),
+ DLDataTypeCode::kDLInt,
+ },
+ {uint32(), DLDataTypeCode::kDLUInt},
+ {
+ int64(),
+ DLDataTypeCode::kDLInt,
+ },
+ {uint64(), DLDataTypeCode::kDLUInt},
+ {float16(), DLDataTypeCode::kDLFloat},
+ {float32(), DLDataTypeCode::kDLFloat},
+ {float64(), DLDataTypeCode::kDLFloat}};
+
+ const auto allocated_bytes = arrow::default_memory_pool()->bytes_allocated();
+
+ for (auto [arrow_type, dlpack_type] : cases) {
+ const std::shared_ptr<Array> array =
+ ArrayFromJSON(arrow_type, "[1, 0, 10, 0, 2, 1, 3, 5, 1, 0]");
+ CheckDLTensor(array, arrow_type, dlpack_type, 10);
+ ASSERT_OK_AND_ASSIGN(auto sliced_1, array->SliceSafe(1, 5));
+ CheckDLTensor(sliced_1, arrow_type, dlpack_type, 5);
+ ASSERT_OK_AND_ASSIGN(auto sliced_2, array->SliceSafe(0, 5));
+ CheckDLTensor(sliced_2, arrow_type, dlpack_type, 5);
+ ASSERT_OK_AND_ASSIGN(auto sliced_3, array->SliceSafe(3));
+ CheckDLTensor(sliced_3, arrow_type, dlpack_type, 7);
+ }
+
+ ASSERT_EQ(allocated_bytes, arrow::default_memory_pool()->bytes_allocated());
+}
+
+TEST_F(TestExportArray, TestErrors) {
+ const std::shared_ptr<Array> array_null = ArrayFromJSON(null(), "[]");
+ ASSERT_RAISES_WITH_MESSAGE(TypeError,
+ "Type error: DataType is not compatible with
DLPack spec: " +
+ array_null->type()->ToString(),
+ arrow::dlpack::ExportArray(array_null));
+
+ const std::shared_ptr<Array> array_with_null = ArrayFromJSON(int8(), "[1,
100, null]");
+ ASSERT_RAISES_WITH_MESSAGE(TypeError,
+ "Type error: Can only use DLPack on arrays with
no nulls.",
+ arrow::dlpack::ExportArray(array_with_null));
+
+ const std::shared_ptr<Array> array_string =
+ ArrayFromJSON(utf8(), R"(["itsy", "bitsy", "spider"])");
+ ASSERT_RAISES_WITH_MESSAGE(TypeError,
+ "Type error: DataType is not compatible with
DLPack spec: " +
+ array_string->type()->ToString(),
+ arrow::dlpack::ExportArray(array_string));
+
+ const std::shared_ptr<Array> array_boolean = ArrayFromJSON(boolean(),
"[true, false]");
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError, "Type error: Bit-packed boolean data type not supported by
DLPack.",
+ arrow::dlpack::ExportDevice(array_boolean));
+}
+
+} // namespace arrow::dlpack
diff --git a/dev/release/rat_exclude_files.txt
b/dev/release/rat_exclude_files.txt
index ce637bf839..4f86a12afe 100644
--- a/dev/release/rat_exclude_files.txt
+++ b/dev/release/rat_exclude_files.txt
@@ -12,6 +12,7 @@ ci/etc/*.patch
ci/vcpkg/*.patch
CHANGELOG.md
cpp/CHANGELOG_PARQUET.md
+cpp/src/arrow/c/dlpack_abi.h
cpp/src/arrow/io/mman.h
cpp/src/arrow/util/random.h
cpp/src/arrow/status.cc
diff --git a/docs/source/python/dlpack.rst b/docs/source/python/dlpack.rst
new file mode 100644
index 0000000000..f612ebabde
--- /dev/null
+++ b/docs/source/python/dlpack.rst
@@ -0,0 +1,93 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+.. or more contributor license agreements. See the NOTICE file
+.. distributed with this work for additional information
+.. regarding copyright ownership. The ASF licenses this file
+.. to you under the Apache License, Version 2.0 (the
+.. "License"); you may not use this file except in compliance
+.. with the License. You may obtain a copy of the License at
+
+.. http://www.apache.org/licenses/LICENSE-2.0
+
+.. Unless required by applicable law or agreed to in writing,
+.. software distributed under the License is distributed on an
+.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+.. KIND, either express or implied. See the License for the
+.. specific language governing permissions and limitations
+.. under the License.
+
+.. _pyarrow-dlpack:
+
+The DLPack Protocol
+===================
+
+`The DLPack Protocol <https://github.com/dmlc/dlpack>`_
+is a stable in-memory data structure that allows exchange
+between major frameworks working with multidimensional
+arrays or tensors. It is designed for cross hardware
+support meaning it allows exchange of data on devices other
+than the CPU (e.g. GPU).
+
+DLPack protocol had been
+`selected as the Python array API standard
<https://data-apis.org/array-api/latest/design_topics/data_interchange.html#dlpack-an-in-memory-tensor-structure>`_
+by the
+`Consortium for Python Data API Standards <https://data-apis.org/>`_
+in order to enable device aware data interchange between array/tensor
+libraries in the Python ecosystem. See more about the standard
+in the
+`protocol documentation <https://data-apis.org/array-api/latest/index.html>`_
+and more about DLPack in the
+`Python Specification for DLPack
<https://dmlc.github.io/dlpack/latest/python_spec.html#python-spec>`_.
+
+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
+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
+ memory.
+
+2. ``__dlpack__(self, stream=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``.
+
+Example
+-------
+
+Convert a PyArrow CPU array to NumPy array:
+
+.. code-block::
+
+ >>> import pyarrow as pa
+ >>> array = pa.array([2, 0, 2, 4])
+ <pyarrow.lib.Int64Array object at 0x121fd4880>
+ [
+ 2,
+ 0,
+ 2,
+ 4
+ ]
+
+ >>> import numpy as np
+ >>> np.from_dlpack(array)
+ array([2, 0, 2, 4])
+
+Convert a PyArrow CPU array to PyTorch tensor:
+
+.. code-block::
+
+ >>> import torch
+ >>> torch.from_dlpack(array)
+ tensor([2, 0, 2, 4])
diff --git a/docs/source/python/index.rst b/docs/source/python/index.rst
index 6a3de3d42b..08939bc760 100644
--- a/docs/source/python/index.rst
+++ b/docs/source/python/index.rst
@@ -53,6 +53,7 @@ files into Arrow structures.
numpy
pandas
interchange_protocol
+ dlpack
timestamps
orc
csv
diff --git a/docs/source/python/interchange_protocol.rst
b/docs/source/python/interchange_protocol.rst
index c354541a67..2a5ec8afed 100644
--- a/docs/source/python/interchange_protocol.rst
+++ b/docs/source/python/interchange_protocol.rst
@@ -37,7 +37,7 @@ libraries in the Python ecosystem. See more about the
standard in the
`protocol documentation
<https://data-apis.org/dataframe-protocol/latest/index.html>`_.
-From pyarrow to other libraries: ``__dataframe__()`` method
+From PyArrow to other libraries: ``__dataframe__()`` method
-----------------------------------------------------------
The ``__dataframe__()`` method creates a new exchange object that
@@ -54,7 +54,7 @@ This is meant to be used by the consumer library when calling
the ``from_dataframe()`` function and is not meant to be used manually
by the user.
-From other libraries to pyarrow: ``from_dataframe()``
+From other libraries to PyArrow: ``from_dataframe()``
-----------------------------------------------------
With the ``from_dataframe()`` function, we can construct a
:class:`pyarrow.Table`
@@ -63,7 +63,7 @@ from any dataframe object that implements the
protocol.
We can for example take a pandas dataframe and construct a
-pyarrow table with the use of the interchange protocol:
+PyArrow table with the use of the interchange protocol:
.. code-block::
diff --git a/python/pyarrow/_dlpack.pxi b/python/pyarrow/_dlpack.pxi
new file mode 100644
index 0000000000..c2f4cff640
--- /dev/null
+++ b/python/pyarrow/_dlpack.pxi
@@ -0,0 +1,46 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+cimport cpython
+from cpython.pycapsule cimport PyCapsule_New
+
+
+cdef void dlpack_pycapsule_deleter(object dltensor) noexcept:
+ cdef DLManagedTensor* dlm_tensor
+ cdef PyObject* err_type
+ cdef PyObject* err_value
+ cdef PyObject* err_traceback
+
+ # Do nothing if the capsule has been consumed
+ if cpython.PyCapsule_IsValid(dltensor, "used_dltensor"):
+ return
+
+ # An exception may be in-flight, we must save it in case
+ # we create another one
+ cpython.PyErr_Fetch(&err_type, &err_value, &err_traceback)
+
+ dlm_tensor = <DLManagedTensor*>cpython.PyCapsule_GetPointer(dltensor,
'dltensor')
+ if dlm_tensor == NULL:
+ cpython.PyErr_WriteUnraisable(dltensor)
+ # The deleter can be NULL if there is no way for the caller
+ # to provide a reasonable destructor
+ elif dlm_tensor.deleter:
+ dlm_tensor.deleter(dlm_tensor)
+ assert (not cpython.PyErr_Occurred())
+
+ # Set the error indicator from err_type, err_value, err_traceback
+ cpython.PyErr_Restore(err_type, err_value, err_traceback)
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index 789e30d3e9..74a196002b 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -1779,6 +1779,44 @@ cdef class Array(_PandasConvertible):
return pyarrow_wrap_array(array)
+ def __dlpack__(self, stream=None):
+ """Export a primitive array as a DLPack capsule.
+
+ Parameters
+ ----------
+ stream : int, optional
+ A Python integer representing a pointer to a stream. Currently not
supported.
+ Stream is provided by the consumer to the producer to instruct the
producer
+ to ensure that operations can safely be performed on the array.
+
+ Returns
+ -------
+ capsule : PyCapsule
+ A DLPack capsule for the array, pointing to a DLManagedTensor.
+ """
+ if stream is None:
+ dlm_tensor = GetResultValue(ExportToDLPack(self.sp_array))
+
+ return PyCapsule_New(dlm_tensor, 'dltensor',
dlpack_pycapsule_deleter)
+ else:
+ raise NotImplementedError(
+ "Only stream=None is supported."
+ )
+
+ def __dlpack_device__(self):
+ """
+ Return the DLPack device tuple this arrays resides on.
+
+ Returns
+ -------
+ tuple : Tuple[int, int]
+ Tuple with index specifying the type of the device (where
+ CPU = 1, see cpp/src/arrow/c/dpack_abi.h) and index of the
+ device which is 0 by default for CPU.
+ """
+ device = GetResultValue(ExportDevice(self.sp_array))
+ return device.device_type, device.device_id
+
cdef _array_like_to_pandas(obj, options, types_mapper):
cdef:
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index 403846a38f..bad5ec606c 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -1199,6 +1199,25 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
shared_ptr[CScalar] MakeNullScalar(shared_ptr[CDataType] type)
+cdef extern from "arrow/c/dlpack_abi.h" nogil:
+ ctypedef enum DLDeviceType:
+ kDLCPU = 1
+
+ ctypedef struct DLDevice:
+ DLDeviceType device_type
+ int32_t device_id
+
+ ctypedef struct DLManagedTensor:
+ void (*deleter)(DLManagedTensor*)
+
+
+cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil:
+ CResult[DLManagedTensor*] ExportToDLPack" arrow::dlpack::ExportArray"(
+ const shared_ptr[CArray]& arr)
+
+ CResult[DLDevice] ExportDevice(const shared_ptr[CArray]& arr)
+
+
cdef extern from "arrow/builder.h" namespace "arrow" nogil:
cdef cppclass CArrayBuilder" arrow::ArrayBuilder":
diff --git a/python/pyarrow/lib.pyx b/python/pyarrow/lib.pyx
index 57fb0f42e3..29a0bed559 100644
--- a/python/pyarrow/lib.pyx
+++ b/python/pyarrow/lib.pyx
@@ -176,6 +176,9 @@ include "table.pxi"
# Tensors
include "tensor.pxi"
+# DLPack
+include "_dlpack.pxi"
+
# File IO
include "io.pxi"
diff --git a/python/pyarrow/tests/test_dlpack.py
b/python/pyarrow/tests/test_dlpack.py
new file mode 100644
index 0000000000..7cf3f4acdb
--- /dev/null
+++ b/python/pyarrow/tests/test_dlpack.py
@@ -0,0 +1,142 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import ctypes
+from functools import wraps
+import pytest
+
+import numpy as np
+
+import pyarrow as pa
+from pyarrow.vendored.version import Version
+
+
+def PyCapsule_IsValid(capsule, name):
+ return ctypes.pythonapi.PyCapsule_IsValid(ctypes.py_object(capsule), name)
== 1
+
+
+def check_dlpack_export(arr, expected_arr):
+ DLTensor = arr.__dlpack__()
+ assert PyCapsule_IsValid(DLTensor, b"dltensor") is True
+
+ result = np.from_dlpack(arr)
+ np.testing.assert_array_equal(result, expected_arr, strict=True)
+
+ assert arr.__dlpack_device__() == (1, 0)
+
+
+def check_bytes_allocated(f):
+ @wraps(f)
+ def wrapper(*args, **kwargs):
+ allocated_bytes = pa.total_allocated_bytes()
+ try:
+ return f(*args, **kwargs)
+ finally:
+ assert pa.total_allocated_bytes() == allocated_bytes
+ return wrapper
+
+
+@check_bytes_allocated
[email protected](
+ ('value_type', 'np_type'),
+ [
+ (pa.uint8(), np.uint8),
+ (pa.uint16(), np.uint16),
+ (pa.uint32(), np.uint32),
+ (pa.uint64(), np.uint64),
+ (pa.int8(), np.int8),
+ (pa.int16(), np.int16),
+ (pa.int32(), np.int32),
+ (pa.int64(), np.int64),
+ (pa.float16(), np.float16),
+ (pa.float32(), np.float32),
+ (pa.float64(), np.float64),
+ ]
+)
+def test_dlpack(value_type, np_type):
+ if Version(np.__version__) < Version("1.24.0"):
+ pytest.skip("No dlpack support in numpy versions older than 1.22.0, "
+ "strict keyword in assert_array_equal added in numpy
version "
+ "1.24.0")
+
+ expected = np.array([1, 2, 3], dtype=np_type)
+ arr = pa.array(expected, type=value_type)
+ check_dlpack_export(arr, expected)
+
+ arr_sliced = arr.slice(1, 1)
+ expected = np.array([2], dtype=np_type)
+ check_dlpack_export(arr_sliced, expected)
+
+ arr_sliced = arr.slice(0, 1)
+ expected = np.array([1], dtype=np_type)
+ check_dlpack_export(arr_sliced, expected)
+
+ arr_sliced = arr.slice(1)
+ expected = np.array([2, 3], dtype=np_type)
+ check_dlpack_export(arr_sliced, expected)
+
+ arr_zero = pa.array([], type=value_type)
+ expected = np.array([], dtype=np_type)
+ check_dlpack_export(arr_zero, expected)
+
+
+def test_dlpack_not_supported():
+ if Version(np.__version__) < Version("1.22.0"):
+ pytest.skip("No dlpack support in numpy versions older than 1.22.0.")
+
+ arr = pa.array([1, None, 3])
+ with pytest.raises(TypeError, match="Can only use DLPack "
+ "on arrays with no nulls."):
+ np.from_dlpack(arr)
+
+ arr = pa.array(
+ [[0, 1], [3, 4]],
+ type=pa.list_(pa.int32())
+ )
+ with pytest.raises(TypeError, match="DataType is not compatible with
DLPack spec"):
+ np.from_dlpack(arr)
+
+ arr = pa.array([])
+ with pytest.raises(TypeError, match="DataType is not compatible with
DLPack spec"):
+ np.from_dlpack(arr)
+
+ # DLPack doesn't support bit-packed boolean values
+ arr = pa.array([True, False, True])
+ with pytest.raises(TypeError, match="Bit-packed boolean data type "
+ "not supported by DLPack."):
+ np.from_dlpack(arr)
+
+
+def test_dlpack_cuda_not_supported():
+ cuda = pytest.importorskip("pyarrow.cuda")
+
+ schema = pa.schema([pa.field('f0', pa.int16())])
+ a0 = pa.array([1, 2, 3], type=pa.int16())
+ batch = pa.record_batch([a0], schema=schema)
+
+ cbuf = cuda.serialize_record_batch(batch, cuda.Context(0))
+ cbatch = cuda.read_record_batch(cbuf, batch.schema)
+ carr = cbatch["f0"]
+
+ # CudaBuffers not yet supported
+ with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
+ "only for buffers on CPU device."):
+ np.from_dlpack(carr)
+
+ with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
+ "only for buffers on CPU device."):
+ carr.__dlpack_device__()
