From bb3453864f7102eb514ea94ef9801143f00f5307 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 4 Sep 2026 13:57:09 +0200 Subject: [PATCH 1/5] Add FixedShapeTensorArray.__dlpack__ --- python/pyarrow/array.pxi | 16 +++++++++++++++- python/pyarrow/tests/test_dlpack.py | 23 +++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 2b2130e992e..691623b87f0 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/tests/test_dlpack.py b/python/pyarrow/tests/test_dlpack.py index f9aac892ced..399ea1cb756 100644 --- a/python/pyarrow/tests/test_dlpack.py +++ b/python/pyarrow/tests/test_dlpack.py @@ -165,6 +165,29 @@ def test_array_to_tensor_dlpack(arr, expected): assert tensor.__dlpack_device__() == (1, 0) +@check_bytes_allocated +def test_fixed_shape_tensor_dlpack_permuted(): + if Version(np.__version__) < Version("2.1.0"): + pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later") + + # 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) + + 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 From 29bbc293b9b6ac1e473cca63ccf85629838eadd4 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 4 Sep 2026 13:57:43 +0200 Subject: [PATCH 2/5] Factor numpy skips --- python/pyarrow/tests/test_dlpack.py | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/python/pyarrow/tests/test_dlpack.py b/python/pyarrow/tests/test_dlpack.py index 399ea1cb756..6ffc776aa98 100644 --- a/python/pyarrow/tests/test_dlpack.py +++ b/python/pyarrow/tests/test_dlpack.py @@ -29,6 +29,13 @@ 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,11 +170,9 @@ 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_dlpack_permuted(): - if Version(np.__version__) < Version("2.1.0"): - pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later") - # 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( @@ -206,12 +209,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() @@ -287,12 +288,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( @@ -300,12 +299,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 From 6790bce8139590d5a69ca997c0c7b3c72c957851 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 4 Sep 2026 14:44:02 +0200 Subject: [PATCH 3/5] Add Dlpack on FixedShapeTensorScalar --- python/pyarrow/scalar.pxi | 29 +++++++++++++++++++++++++++++ python/pyarrow/tests/test_dlpack.py | 27 +++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index fb7de926edc..91be03c7c3d 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 array 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 = self.wrapped.get() + CBaseListScalar* storage = 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 6ffc776aa98..b00ffc0e1c7 100644 --- a/python/pyarrow/tests/test_dlpack.py +++ b/python/pyarrow/tests/test_dlpack.py @@ -191,6 +191,33 @@ def test_fixed_shape_tensor_dlpack_permuted(): 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) + + +@check_bytes_allocated +def test_fixed_shape_tensor_scalar_dlpack_device_permuted(): + # A permuted scalar cannot currently be exported as a Tensor, but the + # device query still resolves through the storage array. + 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) + + assert arr[2].__dlpack_device__() == (1, 0) + + 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 From 6bf37d26277fe14238c88e1423b423fc92e882da Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 4 Sep 2026 14:56:05 +0200 Subject: [PATCH 4/5] Clean up tests --- python/pyarrow/tests/test_dlpack.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/python/pyarrow/tests/test_dlpack.py b/python/pyarrow/tests/test_dlpack.py index b00ffc0e1c7..e3cc2fd3e9e 100644 --- a/python/pyarrow/tests/test_dlpack.py +++ b/python/pyarrow/tests/test_dlpack.py @@ -172,7 +172,7 @@ def test_array_to_tensor_dlpack(arr, expected): @requires_numpy_version("2.1.0") @check_bytes_allocated -def test_fixed_shape_tensor_dlpack_permuted(): +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( @@ -206,18 +206,6 @@ def test_fixed_shape_tensor_scalar_dlpack(): np.testing.assert_array_equal(result, np_arr[1], strict=True) -@check_bytes_allocated -def test_fixed_shape_tensor_scalar_dlpack_device_permuted(): - # A permuted scalar cannot currently be exported as a Tensor, but the - # device query still resolves through the storage array. - 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) - - assert arr[2].__dlpack_device__() == (1, 0) - - 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 From ce5c25aa949fea61e6893eae69c3dac029f99eed Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 4 Sep 2026 14:58:37 +0200 Subject: [PATCH 5/5] Fix typo --- python/pyarrow/scalar.pxi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index 91be03c7c3d..863ab2b66a8 100644 --- a/python/pyarrow/scalar.pxi +++ b/python/pyarrow/scalar.pxi @@ -1588,7 +1588,7 @@ cdef class FixedShapeTensorScalar(ExtensionScalar): def __dlpack__(self, *, stream=None, max_version=None, dl_device=None, copy=None): """ - Export a tensor array as a DLPack capsule. + Export a tensor scalar as a DLPack capsule. See :meth:`Tensor.__dlpack__` for the parameter semantics. """