apacheGH-39294: [C++][Python] DLPack on Tensor class (apache#42118)

AlenkaF · rok · web-flow · commit baa538706fdb · 2025-06-18T11:06:24.000+02:00
### Rationale for this change Producing part of DLPack protocol [has been added to Arrow Arrays](apache#33984) but is missing on Arrow Tensor class. ### What changes are included in this PR? This PR adds support for producing DLPack struct and bindings to it in Python: - `ExportTensor` and `ExportTensorDevice` methods on Arrow C++ Tensor class - `__dlpack__` method on PyArrow Tensor class exposing ExportTensor method - `__dlpack_device__` method on PyArrow Tensor class exposing ExportTensorDevice method ### Are these changes tested? Yes ### Are there any user-facing changes? No * GitHub Issue: apache#39294 Lead-authored-by: AlenkaF <frim.alenka@gmail.com> Co-authored-by: Alenka Frim <AlenkaF@users.noreply.github.com> Co-authored-by: Rok Mihevc <rok@mihevc.org> Signed-off-by: AlenkaF <frim.alenka@gmail.com>
diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc
@@ -20,6 +20,7 @@
 #include "arrow/array/array_base.h"
 #include "arrow/c/dlpack_abi.h"
 #include "arrow/device.h"
+#include "arrow/tensor.h"
 #include "arrow/type.h"
 #include "arrow/type_traits.h"
 
@@ -66,7 +67,7 @@ struct ManagerCtx {
 }  // namespace
 
 Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr) {
-  // Define DLDevice struct nad check if array type is supported
+  // Define DLDevice struct and 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))
@@ -77,7 +78,7 @@ Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr) {
   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);
+  auto ctx = std::make_unique<ManagerCtx>();
 
   // Define the data pointer to the DLTensor
   // If array is of length 0, data pointer should be NULL
@@ -130,4 +131,71 @@ Result<DLDevice> ExportDevice(const std::shared_ptr<Array>& arr) {
   }
 }
 
+struct TensorManagerCtx {
+  std::shared_ptr<Tensor> t;
+  std::vector<int64_t> strides;
+  std::vector<int64_t> shape;
+  DLManagedTensor tensor;
+};
+
+Result<DLManagedTensor*> ExportTensor(const std::shared_ptr<Tensor>& t) {
+  // Define the DLDataType struct
+  const DataType& type = *t->type();
+  ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type));
+
+  // Define DLDevice struct
+  ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(t))
+
+  // Create TensorManagerCtx that will serve as the owner of the DLManagedTensor
+  auto ctx = std::make_unique<TensorManagerCtx>();
+
+  // Define the data pointer to the DLTensor
+  // If tensor is of length 0, data pointer should be NULL
+  if (t->size() == 0) {
+    ctx->tensor.dl_tensor.data = NULL;
+  } else {
+    ctx->tensor.dl_tensor.data = t->raw_mutable_data();
+  }
+
+  ctx->tensor.dl_tensor.device = device;
+  ctx->tensor.dl_tensor.ndim = t->ndim();
+  ctx->tensor.dl_tensor.dtype = dlpack_type;
+  ctx->tensor.dl_tensor.byte_offset = 0;
+
+  std::vector<int64_t>& shape_arr = ctx->shape;
+  shape_arr.reserve(t->ndim());
+  for (auto i : t->shape()) {
+    shape_arr.emplace_back(i);
+  }
+  ctx->tensor.dl_tensor.shape = shape_arr.data();
+
+  std::vector<int64_t>& strides_arr = ctx->strides;
+  strides_arr.reserve(t->ndim());
+  auto byte_width = t->type()->byte_width();
+  for (auto i : t->strides()) {
+    strides_arr.emplace_back(i / byte_width);
+  }
+  ctx->tensor.dl_tensor.strides = strides_arr.data();
+
+  ctx->t = std::move(t);
+  ctx->tensor.manager_ctx = ctx.get();
+  ctx->tensor.deleter = [](struct DLManagedTensor* self) {
+    delete reinterpret_cast<TensorManagerCtx*>(self->manager_ctx);
+  };
+  return &ctx.release()->tensor;
+}
+
+Result<DLDevice> ExportDevice(const std::shared_ptr<Tensor>& t) {
+  // Define DLDevice struct
+  DLDevice device;
+  if (t->data()->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
@@ -39,6 +39,9 @@ namespace arrow::dlpack {
 ARROW_EXPORT
 Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr);
 
+ARROW_EXPORT
+Result<DLManagedTensor*> ExportTensor(const std::shared_ptr<Tensor>& t);
+
 /// \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.
@@ -48,4 +51,7 @@ Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr);
 ARROW_EXPORT
 Result<DLDevice> ExportDevice(const std::shared_ptr<Array>& arr);
 
+ARROW_EXPORT
+Result<DLDevice> ExportDevice(const std::shared_ptr<Tensor>& t);
+
 }  // namespace arrow::dlpack
diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc
@@ -21,6 +21,7 @@
 #include "arrow/c/dlpack.h"
 #include "arrow/c/dlpack_abi.h"
 #include "arrow/memory_pool.h"
+#include "arrow/tensor.h"
 #include "arrow/testing/gtest_util.h"
 
 namespace arrow::dlpack {
@@ -48,7 +49,6 @@ void CheckDLTensor(const std::shared_ptr<Array>& arr,
   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);
@@ -126,4 +126,93 @@ TEST_F(TestExportArray, TestErrors) {
       arrow::dlpack::ExportDevice(array_boolean));
 }
 
+class TestExportTensor : public ::testing::Test {
+ public:
+  void SetUp() {}
+};
+
+void CheckDLTensor(const std::shared_ptr<Tensor>& t,
+                   const std::shared_ptr<DataType>& tensor_type,
+                   DLDataTypeCode dlpack_type, std::vector<int64_t> shape,
+                   std::vector<int64_t> strides) {
+  ASSERT_OK_AND_ASSIGN(auto dlmtensor, arrow::dlpack::ExportTensor(t));
+  auto dltensor = dlmtensor->dl_tensor;
+
+  ASSERT_EQ(t->data()->data(), dltensor.data);
+  ASSERT_EQ(t->ndim(), dltensor.ndim);
+  ASSERT_EQ(0, dltensor.byte_offset);
+  for (int i = 0; i < t->ndim(); i++) {
+    ASSERT_EQ(shape.data()[i], dltensor.shape[i]);
+    ASSERT_EQ(strides.data()[i], dltensor.strides[i]);
+  }
+
+  ASSERT_EQ(dlpack_type, dltensor.dtype.code);
+  ASSERT_EQ(tensor_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(t));
+  ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type);
+  ASSERT_EQ(0, device.device_id);
+
+  dlmtensor->deleter(dlmtensor);
+}
+
+TEST_F(TestExportTensor, TestTensor) {
+  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) {
+    std::vector<int64_t> shape = {3, 6};
+    std::vector<int64_t> dlpack_strides = {6, 1};
+    std::shared_ptr<Tensor> tensor = TensorFromJSON(
+        arrow_type, "[1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9]", shape);
+
+    CheckDLTensor(tensor, arrow_type, dlpack_type, shape, dlpack_strides);
+  }
+
+  ASSERT_EQ(allocated_bytes, arrow::default_memory_pool()->bytes_allocated());
+}
+
+TEST_F(TestExportTensor, TestTensorStrided) {
+  std::vector<int64_t> shape = {2, 2, 2};
+  std::vector<int64_t> strides = {sizeof(float) * 4, sizeof(float) * 2,
+                                  sizeof(float) * 1};
+  std::vector<int64_t> dlpack_strides = {4, 2, 1};
+  std::shared_ptr<Tensor> tensor =
+      TensorFromJSON(float32(), "[1, 2, 3, 4, 5, 6, 1, 1]", shape, strides);
+
+  CheckDLTensor(tensor, float32(), DLDataTypeCode::kDLFloat, shape, dlpack_strides);
+
+  std::vector<int64_t> f_strides = {sizeof(float) * 1, sizeof(float) * 2,
+                                    sizeof(float) * 4};
+  std::vector<int64_t> f_dlpack_strides = {1, 2, 4};
+  std::shared_ptr<Tensor> f_tensor =
+      TensorFromJSON(float32(), "[1, 2, 3, 4, 5, 6, 1, 1]", shape, f_strides);
+
+  CheckDLTensor(f_tensor, float32(), DLDataTypeCode::kDLFloat, shape, f_dlpack_strides);
+}
+
 }  // namespace arrow::dlpack
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
@@ -2144,7 +2144,8 @@ cdef class Array(_PandasConvertible):
         return pyarrow_wrap_array(array)
 
     def __dlpack__(self, stream=None):
-        """Export a primitive array as a DLPack capsule.
+        """
+        Export a primitive array as a DLPack capsule.
 
         Parameters
         ----------
@@ -2159,7 +2160,7 @@ cdef class Array(_PandasConvertible):
             A DLPack capsule for the array, pointing to a DLManagedTensor.
         """
         if stream is None:
-            dlm_tensor = GetResultValue(ExportToDLPack(self.sp_array))
+            dlm_tensor = GetResultValue(ExportArrayToDLPack(self.sp_array))
 
             return PyCapsule_New(dlm_tensor, 'dltensor', dlpack_pycapsule_deleter)
         else:
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
@@ -1452,10 +1452,13 @@ cdef extern from "arrow/c/dlpack_abi.h" nogil:
 
 
 cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil:
-    CResult[DLManagedTensor*] ExportToDLPack" arrow::dlpack::ExportArray"(
+    CResult[DLManagedTensor*] ExportArrayToDLPack" arrow::dlpack::ExportArray"(
         const shared_ptr[CArray]& arr)
+    CResult[DLManagedTensor*] ExportTensorToDLPack" arrow::dlpack::ExportTensor"(
+        const shared_ptr[CTensor]& tensor)
 
     CResult[DLDevice] ExportDevice(const shared_ptr[CArray]& arr)
+    CResult[DLDevice] ExportDevice(const shared_ptr[CTensor]& tensor)
 
 
 cdef extern from "arrow/builder.h" namespace "arrow" nogil:
diff --git a/python/pyarrow/tensor.pxi b/python/pyarrow/tensor.pxi
@@ -300,6 +300,45 @@ strides: {self.strides}"""
         buffer.strides = <Py_ssize_t *> cp.PyBytes_AsString(self._ssize_t_strides)
         buffer.suboffsets = NULL
 
+    def __dlpack__(self, stream=None):
+        """
+        Export a Tensor 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 tensor, pointing to a DLManagedTensor.
+        """
+        if stream is None:
+            dlm_tensor = GetResultValue(ExportTensorToDLPack(self.sp_tensor))
+
+            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 tensor 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_tensor))
+        return device.device_type, device.device_id
+
 
 ctypedef CSparseCOOIndex* _CSparseCOOIndexPtr
 
diff --git a/python/pyarrow/tests/test_dlpack.py b/python/pyarrow/tests/test_dlpack.py
@@ -20,18 +20,13 @@
 import gc
 import pytest
 
-try:
-    import numpy as np
-except ImportError:
-    np = None
-
 import pyarrow as pa
 from pyarrow.vendored.version import Version
 
-
 # Marks all of the tests in this module
 # Ignore these with pytest ... -m 'not numpy'
 pytestmark = pytest.mark.numpy
+np = pytest.importorskip("numpy")
 
 
 def PyCapsule_IsValid(capsule, name):
@@ -87,6 +82,9 @@ def test_dlpack(value_type, np_type_str):
     arr = pa.array(expected, type=value_type)
     check_dlpack_export(arr, expected)
 
+    t = pa.Tensor.from_numpy(expected)
+    check_dlpack_export(t, expected)
+
     arr_sliced = arr.slice(1, 1)
     expected = np.array([2], dtype=np.dtype(np_type_str))
     check_dlpack_export(arr_sliced, expected)
@@ -103,6 +101,30 @@ def test_dlpack(value_type, np_type_str):
     expected = np.array([], dtype=np.dtype(np_type_str))
     check_dlpack_export(arr_zero, expected)
 
+    t = pa.Tensor.from_numpy(expected)
+    check_dlpack_export(t, expected)
+
+
+@check_bytes_allocated
+@pytest.mark.parametrize('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_dlpack(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")
+
+    arr = np.array([1, 2, 3, 4, 5, 6, 1, 1])
+    expected = np.array(arr, dtype=np_type).reshape((2, 2, 2), order='C')
+    t = pa.Tensor.from_numpy(expected)
+    check_dlpack_export(t, expected)
+
+    expected = np.array(arr, dtype=np_type).reshape((2, 2, 2), order='F')
+    t = pa.Tensor.from_numpy(expected)
+    check_dlpack_export(t, expected)
+
 
 def test_dlpack_not_supported():
     if Version(np.__version__) < Version("1.22.0"):
