GH-38868: [Python] Add dlpack producer to FixedShapeTensorArray/Scalar

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):
     """

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):
     """

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,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,25 +303,21 @@ 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(
             DLPackForwarder(obj, max_version=(1, 0), copy=copy))
         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