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 9808d818ce GH-38868: [Python] Add dlpack producer to
FixedShapeTensorArray/Scalar (#51159)
9808d818ce is described below
commit 9808d818ce4ce3321394c647149de849168cdfaa
Author: Antoine Prouvost <[email protected]>
AuthorDate: Mon Sep 7 10:57:21 2026 +0200
GH-38868: [Python] Add dlpack producer to FixedShapeTensorArray/Scalar
(#51159)
### Rationale for this change
Have the same coverage as `to_numpy`.
### What changes are included in this PR?
- `FixedShapeTensorArray.__dlpack__` explicitly calling `to_tensor`
- `FixedShapeTensorArray.__dlpack_device__` (defaulted)
- `FixedShapeTensorScalar.__dlpack__` explicitly calling `to_tensor`
- `FixedShapeTensorScalar.__dlpack_device__` calling C++ device extraction
on the underlying array.
If GH-51122 get merged first, I'll add the consuming side here, otherwise
if this one gets merged first, I'll add it there.
### Are these changes tested?
Yes a few since this does not introduce new logic.
### Are there any user-facing changes?
Additions only.
* GitHub Issue: #38868
Authored-by: AntoinePrv <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
python/pyarrow/array.pxi | 16 +++++++++-
python/pyarrow/scalar.pxi | 29 ++++++++++++++++++
python/pyarrow/tests/test_dlpack.py | 59 +++++++++++++++++++++++++++++--------
3 files changed, 91 insertions(+), 13 deletions(-)
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index 2b2130e992..691623b87f 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -2269,7 +2269,7 @@ cdef class Array(_PandasConvertible):
return pyarrow_wrap_array(array)
- def __dlpack__(self, stream=None, max_version=None, dl_device=None,
copy=None):
+ def __dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None):
"""
Export a primitive array as a DLPack capsule.
@@ -5045,6 +5045,20 @@ cdef class FixedShapeTensorArray(ExtensionArray):
FixedSizeListArray.from_arrays(values, shape[1:].prod())
)
+ def __dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None):
+ """
+ Export a tensor array as a DLPack capsule.
+
+ The element positions in the array become the first dimension of the
+ resulting tensor (equal to ``len(self)``).
+
+ See :meth:`Tensor.__dlpack__` for the parameter semantics.
+ """
+ return self.to_tensor().__dlpack__(
+ stream=stream, max_version=max_version,
+ dl_device=dl_device, copy=copy,
+ )
+
cdef class OpaqueArray(ExtensionArray):
"""
diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi
index fb7de926ed..863ab2b66a 100644
--- a/python/pyarrow/scalar.pxi
+++ b/python/pyarrow/scalar.pxi
@@ -1586,6 +1586,35 @@ cdef class FixedShapeTensorScalar(ExtensionScalar):
ctensor = GetResultValue(c_type.MakeTensor(scalar))
return pyarrow_wrap_tensor(ctensor)
+ def __dlpack__(self, *, stream=None, max_version=None, dl_device=None,
copy=None):
+ """
+ Export a tensor scalar as a DLPack capsule.
+
+ See :meth:`Tensor.__dlpack__` for the parameter semantics.
+ """
+ return self.to_tensor().__dlpack__(
+ stream=stream, max_version=max_version,
+ dl_device=dl_device, copy=copy,
+ )
+
+ def __dlpack_device__(self):
+ """
+ Return the DLPack device tuple this scalar resides on.
+
+ Returns
+ -------
+ tuple : Tuple[int, int]
+ Tuple with index specifying the type of the device (where
+ CPU = 1, see cpp/src/arrow/c/dlpack_abi.h) and index of the
+ device which is 0 by default for CPU.
+ """
+ cdef:
+ CExtensionScalar* ext = <CExtensionScalar*> self.wrapped.get()
+ CBaseListScalar* storage = <CBaseListScalar*> ext.value.get()
+ # The base storage for this type is an Array, so we call into this
function
+ device = GetResultValue(ExportDevice(storage.value))
+ return device.device_type, device.device_id
+
cdef class OpaqueScalar(ExtensionScalar):
"""
diff --git a/python/pyarrow/tests/test_dlpack.py
b/python/pyarrow/tests/test_dlpack.py
index f9aac892ce..e3cc2fd3e9 100644
--- a/python/pyarrow/tests/test_dlpack.py
+++ b/python/pyarrow/tests/test_dlpack.py
@@ -29,6 +29,13 @@ pytestmark = pytest.mark.numpy
np = pytest.importorskip("numpy")
+def requires_numpy_version(min_version):
+ return pytest.mark.skipif(
+ Version(np.__version__) < Version(min_version),
+ reason=f"Test requires numpy {min_version} or later",
+ )
+
+
def PyCapsule_IsValid(capsule, name):
return ctypes.pythonapi.PyCapsule_IsValid(ctypes.py_object(capsule), name)
== 1
@@ -150,12 +157,10 @@ def multidim_arrays():
]
+@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize(('arr', 'expected'), multidim_arrays())
def test_array_to_tensor_dlpack(arr, expected):
- if Version(np.__version__) < Version("2.1.0"):
- pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
-
tensor = arr.to_tensor()
# A Tensor sharing an Array buffer is immutable, so it can only be exported
# through the versioned DLPack protocol.
@@ -165,6 +170,42 @@ def test_array_to_tensor_dlpack(arr, expected):
assert tensor.__dlpack_device__() == (1, 0)
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
+def test_fixed_shape_tensor_array_dlpack_permuted():
+ # A non-trivial permutation makes to_tensor() produce a non-row-major
+ # tensor: each row-major [3, 2] block is exposed as a logical [2, 3] cell.
+ storage = pa.FixedSizeListArray.from_arrays(
+ pa.array(range(24), type=pa.int32()), 6)
+ arr = pa.ExtensionArray.from_storage(
+ pa.fixed_shape_tensor(pa.int32(), [3, 2], permutation=[1, 0]), storage)
+
+ tensor = arr.to_tensor()
+ assert tensor.shape == (4, 2, 3)
+ assert not tensor.is_contiguous
+
+ # expected[i, j, k] == i * 6 + k * 2 + j (numpy is only the DLPack
consumer)
+ expected = np.arange(24, dtype=np.int32).reshape(4, 3, 2).transpose(0, 2,
1)
+ result = np.from_dlpack(DLPackForwarder(arr, max_version=(1, 0)))
+ np.testing.assert_array_equal(result, expected, strict=True)
+ assert arr.__dlpack_device__() == (1, 0)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
+def test_fixed_shape_tensor_scalar_dlpack():
+ np_arr = np.arange(12, dtype=np.int32).reshape(3, 2, 2)
+ arr = pa.FixedShapeTensorArray.from_numpy_ndarray(np_arr)
+
+ scalar = arr[1]
+ assert isinstance(scalar, pa.FixedShapeTensorScalar)
+ # __dlpack_device__ reads the storage array's device, without building a
Tensor.
+ assert scalar.__dlpack_device__() == (1, 0)
+
+ result = np.from_dlpack(DLPackForwarder(scalar, max_version=(1, 0)))
+ np.testing.assert_array_equal(result, np_arr[1], strict=True)
+
+
def multidim_arrays_with_nulls():
np_arr = np.arange(6, dtype=np.int32).reshape(3, 2)
# Masked entries keep defined values in the child array, so the tensor
@@ -183,12 +224,10 @@ def multidim_arrays_with_nulls():
]
+@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize(('arr', 'expected'), multidim_arrays_with_nulls())
def test_array_to_tensor_dlpack_nulls(arr, expected):
- if Version(np.__version__) < Version("2.1.0"):
- pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
-
with pytest.raises(pa.ArrowInvalid, match="Array contains nulls"):
arr.to_tensor()
@@ -264,12 +303,10 @@ def test_dlpack_versioned_capsule(obj, max_version, copy):
assert PyCapsule_IsValid(capsule, b"dltensor_versioned") is True
+@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize('obj', dlpack_objects())
def test_dlpack_versioned_roundtrip(obj):
- if Version(np.__version__) < Version("2.1.0"):
- pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
-
expected = np.from_dlpack(DLPackForwarder(obj, max_version=None))
for copy in [None, False, True]:
result = np.from_dlpack(
@@ -277,12 +314,10 @@ def test_dlpack_versioned_roundtrip(obj):
np.testing.assert_array_equal(result, expected, strict=True)
+@requires_numpy_version("2.2.5")
@check_bytes_allocated
def test_dlpack_copy_is_writeable():
# NumPy did not set the writeable flag on DLPack imports before 2.2.5.
- if Version(np.__version__) < Version("2.2.5"):
- pytest.skip("Writable DLPack imports require numpy 2.2.5 or later")
-
arr = pa.array([1, 2, 3], type=pa.int32())
# Arrow arrays are immutable, so a shared export is read-only
