GH-38868: [Python] Add dlpack producer to FixedShapeTensorArray/Scalar - #51159
GH-38868: [Python] Add dlpack producer to FixedShapeTensorArray/Scalar#51159AntoinePrv wants to merge 5 commits into
Conversation
|
Thanks for opening a pull request! This pull request has been automatically converted to a draft because its title doesn't match Arrow's required format. If this is not a minor PR. Could you open an issue for this pull request on GitHub? https://github.com/apache/arrow/issues/new/choose Opening GitHub issues ahead of time contributes to the Openness of the Apache Arrow project. Then could you also rename the pull request title in the following format? or After updating the title, you can mark the pull request as ready for review. See also: |
There was a problem hiding this comment.
🟡 Changes recommended
It introduces a backward-incompatible change to the public Array.__dlpack__ signature and contains a newly added incorrect docstring for FixedShapeTensorScalar.__dlpack__.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR extends PyArrow’s DLPack producer support to FixedShapeTensorArray and FixedShapeTensorScalar, with accompanying tests to match the existing to_numpy-level coverage, and refactors NumPy version gating in DLPack tests.
Changes:
- Add a reusable
requires_numpy_version()marker and apply it to version-dependent DLPack tests. - Implement
__dlpack__forFixedShapeTensorArrayand both__dlpack__/__dlpack_device__forFixedShapeTensorScalar. - Add new DLPack roundtrip tests for permuted fixed-shape tensor arrays and scalars.
File summaries
| File | Description |
|---|---|
| python/pyarrow/tests/test_dlpack.py | Adds NumPy-version gating helper and new tests for fixed-shape tensor array/scalar DLPack export behavior. |
| python/pyarrow/scalar.pxi | Adds DLPack producer/device hooks for FixedShapeTensorScalar. |
| python/pyarrow/array.pxi | Adds DLPack producer for FixedShapeTensorArray and changes Array.__dlpack__ signature. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
🔵 Needs a closer look
FixedShapeTensorScalar.__dlpack_device__ can dereference null storage for invalid scalars, risking a crash instead of raising a Python exception.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
python/pyarrow/scalar.pxi:1616
__dlpack_device__dereferencesext.value/storage.valuewithout checking scalar validity. For a null (invalid)FixedShapeTensorScalarthis can be a null pointer and crash instead of raising a Python exception. Please guard onself.is_validbefore using the storage scalar.
- Files reviewed: 3/3 changed files
- Comments generated: 0 new
- Review effort level: Lite
Rationale for this change
Have the same coverage as
to_numpy.What changes are included in this PR?
FixedShapeTensorArray.__dlpack__explicitly callingto_tensorFixedShapeTensorArray.__dlpack_device__(defaulted)FixedShapeTensorScalar.__dlpack__explicitly callingto_tensorFixedShapeTensorScalar.__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.