You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by pa...@apache.org on 2024/05/24 17:45:36 UTC
(arrow-nanoarrow) branch main updated: feat(python): DLPack support on CBufferView (#469)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new b10621c1 feat(python): DLPack support on CBufferView (#469)
b10621c1 is described below
commit b10621c1f93edeee1db8e5334f5634ad64425902
Author: Alenka Frim <Al...@users.noreply.github.com>
AuthorDate: Fri May 24 19:45:31 2024 +0200
feat(python): DLPack support on CBufferView (#469)
This PR implements producing part of the DLPack protocol on
`CBufferView` (and `CArray`/`CArrayView`) object. With the changes in
this PR, the following code works:
```python
>>> import nanoarrow as na
>>> import numpy as np
>>> pa_arr = pa.array([1, 2, 3], type=pa.uint8())
>>> na_array = na.c_array(pa_arr)
>>> na_view = na_array.view()
>>> view = na_view.buffer(1)
>>> np.from_dlpack(view)
array([1, 2, 3], dtype=uint8)
```
## Changes
Included changes in this PR:
- [x] Implementation of `__dlpack__` method on `CBufferView`
- [x] Implementation of `__dlpack_device__` method on `CBufferView`
- [x] Tests
Possible follow-up
- Add `__dlpack__` and `__dlpack_device__` method to the
`CArray`(`View`?)
- Add implementation for consuming objects that support `__dlpack__`
## About DLPack:
DLPack is selected for Array API protocol. See more about it:
https://dmlc.github.io/dlpack/latest/python_spec.html
---------
Co-authored-by: Dewey Dunnington <de...@dunnington.ca>
---
dev/release/rat_exclude_files.txt | 1 +
python/src/nanoarrow/_lib.pyx | 148 +++++++++++++-
python/src/nanoarrow/dlpack_abi.h | 321 ++++++++++++++++++++++++++++++
python/src/nanoarrow/nanoarrow_dlpack.pxd | 68 +++++++
python/tests/test_dlpack.py | 82 ++++++++
5 files changed, 615 insertions(+), 5 deletions(-)
diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt
index 4bec8dec..14c4153c 100644
--- a/dev/release/rat_exclude_files.txt
+++ b/dev/release/rat_exclude_files.txt
@@ -14,3 +14,4 @@ dist/flatcc/*
dist/flatcc.c
extensions/nanoarrow_ipc/src/nanoarrow/nanoarrow_ipc_flatcc_generated.h
extensions/nanoarrow_ipc/thirdparty/*
+python/src/nanoarrow/dlpack_abi.h
diff --git a/python/src/nanoarrow/_lib.pyx b/python/src/nanoarrow/_lib.pyx
index 590406a0..ce5c36b5 100644
--- a/python/src/nanoarrow/_lib.pyx
+++ b/python/src/nanoarrow/_lib.pyx
@@ -46,11 +46,13 @@ from cpython cimport (
PyBuffer_FillInfo,
PyBUF_ANY_CONTIGUOUS,
PyBUF_FORMAT,
- PyBUF_WRITABLE
+ PyBUF_WRITABLE,
+ PyErr_WriteUnraisable,
)
from cpython.ref cimport Py_INCREF, Py_DECREF
from nanoarrow_c cimport *
from nanoarrow_device_c cimport *
+from nanoarrow_dlpack cimport *
from enum import Enum
from sys import byteorder as sys_byteorder
@@ -166,6 +168,98 @@ cdef object alloc_c_array_view(ArrowArrayView** c_array_view):
return PyCapsule_New(c_array_view[0], 'nanoarrow_array_view', &pycapsule_array_view_deleter)
+cdef void pycapsule_dlpack_deleter(object dltensor) noexcept:
+ cdef DLManagedTensor* dlm_tensor
+
+ # Do nothing if the capsule has been consumed
+ if PyCapsule_IsValid(dltensor, "used_dltensor"):
+ return
+
+ dlm_tensor = <DLManagedTensor*>PyCapsule_GetPointer(dltensor, 'dltensor')
+ if dlm_tensor == NULL:
+ 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)
+
+
+cdef void view_dlpack_deleter(DLManagedTensor* tensor) noexcept with gil:
+ if tensor.manager_ctx is NULL:
+ return
+ Py_DECREF(<CBufferView>tensor.manager_ctx)
+ tensor.manager_ctx = NULL
+ ArrowFree(tensor)
+
+_uint = (NANOARROW_TYPE_UINT8, NANOARROW_TYPE_UINT16, NANOARROW_TYPE_UINT32, NANOARROW_TYPE_UINT64)
+_int = (NANOARROW_TYPE_INT8, NANOARROW_TYPE_INT16, NANOARROW_TYPE_INT32, NANOARROW_TYPE_INT64, NANOARROW_TYPE_INTERVAL_MONTHS)
+_float = (NANOARROW_TYPE_HALF_FLOAT, NANOARROW_TYPE_FLOAT, NANOARROW_TYPE_DOUBLE)
+
+cdef DLDataType view_to_dlpack_data_type(CBufferView view):
+ cdef DLDataType dtype
+ # Define DLDataType struct
+ if view.data_type_id in _uint:
+ dtype.code = kDLUInt
+ elif view.data_type_id in _int:
+ dtype.code = kDLInt
+ elif view.data_type_id in _float:
+ dtype.code = kDLFloat
+ elif view.data_type_id == NANOARROW_TYPE_BOOL:
+ raise ValueError('Bit-packed boolean data type not supported by DLPack.')
+ else:
+ raise ValueError('DataType is not compatible with DLPack spec: ' + view.data_type)
+ dtype.lanes = <uint16_t>1
+ dtype.bits = <uint8_t>(view._element_size_bits)
+
+ return dtype
+
+
+cdef object view_to_dlpack(CBufferView view):
+ # Define DLDevice and DLDataType struct and
+ # with that check for data type support first
+ cdef DLDevice device = view_to_dlpack_device(view)
+ cdef DLDataType dtype = view_to_dlpack_data_type(view)
+
+ # Allocate memory for DLManagedTensor
+ cdef DLManagedTensor* dlm_tensor = <DLManagedTensor*>ArrowMalloc(sizeof(DLManagedTensor))
+ # Define DLManagedTensor struct
+ cdef DLTensor* dl_tensor = &dlm_tensor.dl_tensor
+ dl_tensor.data = <void*>view._ptr.data.data
+ dl_tensor.ndim = 1
+
+ dl_tensor.shape = &view._n_elements
+ dl_tensor.strides = NULL
+ dl_tensor.byte_offset = 0
+
+ dl_tensor.device = device
+ dl_tensor.dtype = dtype
+
+ dlm_tensor.manager_ctx = <void*>view
+ Py_INCREF(view)
+ dlm_tensor.deleter = view_dlpack_deleter
+
+ return PyCapsule_New(dlm_tensor, 'dltensor', pycapsule_dlpack_deleter)
+
+
+cdef DLDevice view_to_dlpack_device(CBufferView view):
+ cdef DLDevice device
+
+ # Check data type support
+ if view.data_type_id == NANOARROW_TYPE_BOOL:
+ raise ValueError('Bit-packed boolean data type not supported by DLPack.')
+ elif view.data_type_id not in _uint + _int + _float:
+ raise ValueError('DataType is not compatible with DLPack spec: ' + view.data_type)
+
+ # Define DLDevice struct
+ if view._device.device_type_id == ARROW_DEVICE_CPU:
+ device.device_type = kDLCPU
+ device.device_id = 0
+ else:
+ raise ValueError('Only CPU device is currently supported.')
+
+ return device
+
+
# Provide a way to validate that we release all references we create
cdef int64_t pyobject_buffer_count = 0
@@ -1842,6 +1936,7 @@ cdef class CBufferView:
cdef Py_ssize_t _element_size_bits
cdef Py_ssize_t _shape
cdef Py_ssize_t _strides
+ cdef int64_t _n_elements
cdef char _format[128]
def __cinit__(self, object base, uintptr_t addr, int64_t size_bytes,
@@ -1861,6 +1956,10 @@ cdef class CBufferView:
)
self._strides = self._item_size()
self._shape = self._ptr.size_bytes // self._strides
+ if self._data_type == NANOARROW_TYPE_BOOL:
+ self._n_elements = self._shape * 8
+ else:
+ self._n_elements = self._shape
def _addr(self):
return <uintptr_t>self._ptr.data.data
@@ -1941,10 +2040,7 @@ cdef class CBufferView:
@property
def n_elements(self):
- if self._data_type == NANOARROW_TYPE_BOOL:
- return self._shape * 8
- else:
- return self._shape
+ return self._n_elements
def element(self, i):
if self._data_type == NANOARROW_TYPE_BOOL:
@@ -2090,6 +2186,48 @@ cdef class CBufferView:
else:
return self._element_size_bits // 8
+
+ def __dlpack__(self, stream=None):
+ """
+ Export CBufferView 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.
+ """
+ # Note: parent offset not applied!
+
+ if stream is None:
+ return view_to_dlpack(self)
+ else:
+ raise NotImplementedError(
+ "Only stream=None is supported."
+ )
+
+
+ def __dlpack_device__(self):
+ """
+ Return the DLPack device tuple this CBufferView resides on.
+
+ Returns
+ -------
+ tuple : Tuple[int, int]
+ Tuple with index specifying the type of the device (where
+ CPU = 1, see python/src/nanoarrow/dpack_abi.h) and index of the
+ device which is 0 by default for CPU.
+ """
+ return (view_to_dlpack_device(self).device_type,
+ view_to_dlpack_device(self).device_id)
+
+
# These are special methods, which can't be cdef and we can't
# call them from elsewhere. We implement the logic for the buffer
# protocol separately so we can re-use it in the CBuffer.
diff --git a/python/src/nanoarrow/dlpack_abi.h b/python/src/nanoarrow/dlpack_abi.h
new file mode 100644
index 00000000..c741c857
--- /dev/null
+++ b/python/src/nanoarrow/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 alignment 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/python/src/nanoarrow/nanoarrow_dlpack.pxd b/python/src/nanoarrow/nanoarrow_dlpack.pxd
new file mode 100644
index 00000000..a9855a68
--- /dev/null
+++ b/python/src/nanoarrow/nanoarrow_dlpack.pxd
@@ -0,0 +1,68 @@
+# 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.
+
+from nanoarrow_c cimport *
+
+cdef extern from "./dlpack_abi.h" nogil:
+ cdef int DLPACK_VERSION
+
+ ctypedef enum DLDeviceType:
+ kDLCPU
+ kDLCUDA
+ kDLCUDAHost
+ kDLOpenCL
+ kDLVulkan
+ kDLMetal
+ kDLVPI
+ kDLROCM
+ kDLROCMHost
+ kDLExtDev
+ kDLCUDAManaged
+ kDLOneAPI
+ kDLWebGPU
+ kDLHexagon
+
+ ctypedef struct DLDevice:
+ DLDeviceType device_type
+ int32_t device_id
+
+ cdef enum DLDataTypeCode:
+ kDLInt
+ kDLUInt
+ kDLFloat
+ kDLBfloat
+ kDLComplex
+ kDLBool
+
+ ctypedef struct DLDataType:
+ uint8_t code
+ uint8_t bits
+ uint16_t lanes
+
+ ctypedef struct DLTensor:
+ void* data
+ DLDevice device
+ int32_t ndim
+ DLDataType dtype
+ int64_t* shape
+ int64_t* strides
+ uint64_t byte_offset
+
+ ctypedef struct DLManagedTensor:
+ DLTensor dl_tensor
+ void* manager_ctx
+ void (*deleter)(DLManagedTensor*) # noqa: E211
diff --git a/python/tests/test_dlpack.py b/python/tests/test_dlpack.py
new file mode 100644
index 00000000..55105763
--- /dev/null
+++ b/python/tests/test_dlpack.py
@@ -0,0 +1,82 @@
+# 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 pytest
+from nanoarrow._lib import _obj_is_capsule
+
+import nanoarrow as na
+
+np = pytest.importorskip("numpy")
+
+
+def check_dlpack_export(view, expected_arr):
+ DLTensor = view.__dlpack__()
+ assert _obj_is_capsule(DLTensor, "dltensor") is True
+
+ result = np.from_dlpack(view)
+ np.testing.assert_array_equal(result, expected_arr, strict=True)
+
+ assert view.__dlpack_device__() == (1, 0)
+
+
+@pytest.mark.parametrize(
+ ("value_type", "np_type"),
+ [
+ (na.uint8(), np.uint8),
+ (na.uint16(), np.uint16),
+ (na.uint32(), np.uint32),
+ (na.uint64(), np.uint64),
+ (na.int8(), np.int8),
+ (na.int16(), np.int16),
+ (na.int32(), np.int32),
+ (na.int64(), np.int64),
+ (na.interval_months(), np.int32),
+ (na.float16(), np.float16),
+ (na.float32(), np.float32),
+ (na.float64(), np.float64),
+ ],
+)
+def test_dlpack(value_type, np_type):
+ if np.__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)
+ # Use the value buffer of the nanoarrow CArray
+ view = na.c_array([1, 2, 3], value_type).view().buffer(1)
+ check_dlpack_export(view, expected)
+
+
+def test_dlpack_not_supported():
+ # DLPack doesn't support bit-packed boolean values
+ view = na.c_array([True, False, True], na.bool_()).view().buffer(1)
+
+ with pytest.raises(
+ ValueError, match="Bit-packed boolean data type not supported by DLPack."
+ ):
+ view.__dlpack__()
+
+ with pytest.raises(
+ ValueError, match="Bit-packed boolean data type not supported by DLPack."
+ ):
+ view.__dlpack_device__()
+
+ with pytest.raises(NotImplementedError, match="Only stream=None is supported."):
+ view.__dlpack__(stream=3)