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/scalar.pxi b/python/pyarrow/scalar.pxi index fb7de926edc..863ab2b66a8 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 = 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 f9aac892ced..e3cc2fd3e9e 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,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