GH-38868: [C++][Python] Add Array::ToTensor and fixed size list support - #50929
GH-38868: [C++][Python] Add Array::ToTensor and fixed size list support#50929AntoinePrv wants to merge 17 commits into
Conversation
|
|
There was a problem hiding this comment.
Pull request overview
This PR introduces a new public Array::ToTensor API (C++ and Python) to enable exporting multidimensional array-like data as Tensor, and updates DLPack export paths and tests to use to_tensor() for multidimensional support (notably nested FixedSizeListArray and FixedShapeTensorArray).
Changes:
- Add virtual
Array::ToTensorplus concrete implementations for 1D numeric arrays and (nested) fixed-size list arrays; routeFixedShapeTensorArray::ToTensorthrough the base virtual. - Refactor tensor stride utilities (row-major stride computation) and simplify DLPack device handling; update DLPack type errors to suggest Tensor conversion.
- Add/extend C++ and Python test coverage for
to_tensor().__dlpack__()on multidimensional inputs.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| python/pyarrow/tests/test_dlpack.py | Adds multidimensional array-to-tensor DLPack export tests via arr.to_tensor() |
| python/pyarrow/includes/libarrow.pxd | Exposes Array::ToTensor() at the Cython API layer |
| python/pyarrow/array.pxi | Adds Array.to_tensor() Python API and routes FixedShapeTensorArray.to_tensor() through it |
| cpp/src/arrow/tensor.h | Updates stride utilities API and adds std::span overload for row-major strides |
| cpp/src/arrow/tensor.cc | Refactors row-major stride computation implementation |
| cpp/src/arrow/extension/fixed_shape_tensor.h | Makes FixedShapeTensorArray::ToTensor() override the new virtual |
| cpp/src/arrow/extension/fixed_shape_tensor.cc | Updates ToTensor() signature to match override |
| cpp/src/arrow/c/dlpack.cc | Refactors DLPack export (device factoring, type checks, array offset/length handling) and updates type errors |
| cpp/src/arrow/c/dlpack_test.cc | Updates DLPack tests to validate shape/strides and revised ExportDevice behavior |
| cpp/src/arrow/array/array_test.cc | Adds C++ unit tests for Array::ToTensor() on primitive arrays |
| cpp/src/arrow/array/array_primitive.h | Implements NumericArray::ToTensor() for 1D numeric arrays |
| cpp/src/arrow/array/array_nested.h | Declares FixedSizeListArray::ToTensor() API |
| cpp/src/arrow/array/array_nested.cc | Implements FixedSizeListArray::ToTensor() with nested fixed-size list support |
| cpp/src/arrow/array/array_list_test.cc | Adds tests for FixedSizeListArray::ToTensor() including nesting, slicing, and null handling |
| cpp/src/arrow/array/array_base.h | Declares new virtual Array::ToTensor() API |
| cpp/src/arrow/array/array_base.cc | Provides default Array::ToTensor() NotImplemented behavior |
💡 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.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 2 comments.
Suppressed comments (1)
python/pyarrow/tests/test_dlpack.py:169
- The numpy version guard checks
< 1.24.0, but the skip message says "No dlpack support ... older than 1.22.0". This is confusing when diagnosing test skips; update the message to reflect the actual minimum version (and optionally mention why 1.24 is required).
if Version(np.__version__) < Version("1.24.0"):
pytest.skip("No dlpack support in numpy versions older than 1.22.0, "
"strict keyword in assert_array_equal added in numpy version "
"1.24.0")
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated no new comments.
Suppressed comments (3)
python/pyarrow/tests/test_dlpack.py:169
- The skip condition is
numpy < 1.24.0, but the message says "older than 1.22.0". This is misleading when diagnosing CI skips; align the message with the actual version gate (or explain both requirements explicitly).
if Version(np.__version__) < Version("1.24.0"):
pytest.skip("No dlpack support in numpy versions older than 1.22.0, "
"strict keyword in assert_array_equal added in numpy version "
"1.24.0")
cpp/src/arrow/array/array_test.cc:1234
- Two of the EXPECT_EQ assertions are no-ops (they compare
shape/stridesto literals that exactly match those variables), so this test isn't actually verifying the tensor shape/strides beyond the later checks. Removing them makes the intent clearer and avoids false confidence in coverage.
EXPECT_EQ(int32(), tensor->type());
EXPECT_EQ(shape, std::vector<int64_t>{5});
EXPECT_EQ(strides, std::vector<int64_t>{sizeof(int32_t)});
EXPECT_EQ(shape, tensor->shape());
EXPECT_EQ(strides, tensor->strides());
cpp/src/arrow/extension/fixed_shape_tensor.h:48
- Docstring grammar: "where this array null entries" is missing a verb. This is a public header comment, so it's worth fixing for clarity.
/// Nulls are ignored, leaving the output tensor with unspecified values where this
/// array null entries.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 2 comments.
Suppressed comments (3)
cpp/src/arrow/array/array_base.h:257
- Public API comment has a couple of grammatical issues ("Example include" / "where this array null entries"), which can be confusing in generated docs.
/// Example include NumericArray, FixedShapeTensorArray, nested FixedSizeListArray.
/// Nulls are ignored, leaving the output tensor with unspecified values where this
/// array null entries.
cpp/src/arrow/array/array_nested.h:653
- Doc comment contains grammatical issues ("number of element", "fixed sized list", "where this array null entries"). Since this is a public override, it will show up in generated docs.
/// The output tensor has a row major layout with the number of element as the first
/// dimension and the fixed sized list as the remaining one (possibly nested).
/// Nulls are ignored, leaving the output tensor with unspecified values where this
/// array null entries.
python/pyarrow/tests/test_dlpack.py:169
- The skip condition is
numpy < 1.24.0, but the message says "No dlpack support ... older than 1.22.0". This is misleading for numpy 1.22/1.23 where dlpack exists but the test still needs 1.24 due tostrict=True.
if Version(np.__version__) < Version("1.24.0"):
pytest.skip("No dlpack support in numpy versions older than 1.22.0, "
"strict keyword in assert_array_equal added in numpy version "
"1.24.0")
dcab253 to
9791a66
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated no new comments.
Suppressed comments (3)
python/pyarrow/tests/test_dlpack.py:169
- The skip condition checks NumPy < 1.24.0, but the message says “older than 1.22.0”. This is confusing when diagnosing CI skips; update the message to match the actual version gate (and/or split the reasons: dlpack support vs
strict=support).
if Version(np.__version__) < Version("1.24.0"):
pytest.skip("No dlpack support in numpy versions older than 1.22.0, "
"strict keyword in assert_array_equal added in numpy version "
"1.24.0")
cpp/src/arrow/c/dlpack.cc:54
- The new TypeError message for unsupported DLPack dtypes always suggests converting to a Tensor, but that advice isn’t applicable to many unsupported types (e.g., strings). Consider qualifying the suggestion (e.g., “if the array represents multi-dimensional numeric data”) or mentioning the concrete API (
Array::ToTensor/to_tensor).
return Status::TypeError("Bit-packed boolean data type not supported by DLPack.");
} else {
return Status::TypeError(
"DataType is not compatible with DLPack spec: ", type.ToString(),
", try converting to a Tensor for multi dimensional data support");
}
cpp/src/arrow/c/dlpack.cc:126
- For the null-count failure path, the error is still correct, but it now misses the new recommended workaround introduced in this PR:
Array::ToTensorsupports nulls (as unspecified values) and can then be exported via DLPack. Consider updating this error message to mention converting to a Tensor as well, for consistency with the new unsupported-type guidance.
Result<DT*> ExportArrayImpl(const std::shared_ptr<Array>& arr, bool copy) {
if (arr->null_count() > 0) {
return Status::TypeError("Can only use DLPack on arrays with no nulls.");
}
|
@AlenkaF this is ready, I think the remaining failures are unrelated. |
|
@rok too :) |
|
Hi @AntoinePrv, will do the initial review today! |
AlenkaF
left a comment
There was a problem hiding this comment.
Went through the PR, amazing work!
I only have two questions regarding the tensor files, other is a nit.
|
|
||
| if (remaining == 0) { | ||
| strides->assign(shape.size(), byte_width); | ||
| // An empty dimension makes the whole tensor empty, so any stride is as good. |
There was a problem hiding this comment.
The changes in tensor.h/.cc were not clear to me so I helped myself with Claude to better understand. Is the reason behind the change (order and code change) meant to catch an empty dimension up front?
| /// Pass `elem_size=1` to get the strides in number of elements, or the element size in | ||
| /// bytes to get them in bytes. On error, the contents of `strides` are unspecified. | ||
| ARROW_EXPORT | ||
| Status ComputeRowMajorStrides(std::span<const int64_t> shape, int64_t elem_size, |
There was a problem hiding this comment.
Is there a specific reason why the API changed (new signature added, the old becoming a thin wrapper)? Is this meant to serve a purpose for DLPack or is this a general improvement? cc @rok
There was a problem hiding this comment.
That is a leftover from when I initially implemented tensor support for fixed size list directly in DLPack (so I needed to compute the strides internally).
Reverting now.
| /// Nulls are ignored, leaving the output tensor with unspecified values where this | ||
| /// array has null entries. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.
Suppressed comments (1)
python/pyarrow/tests/test_dlpack.py:169
- The skip condition checks for NumPy < 1.24.0, but the message says "older than 1.22.0". This is misleading when diagnosing CI skips; update the message to match the actual version gate (or adjust the gate if 1.22 is really sufficient).
if Version(np.__version__) < Version("1.24.0"):
pytest.skip("No dlpack support in numpy versions older than 1.22.0, "
"strict keyword in assert_array_equal added in numpy version "
"1.24.0")
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.
Suppressed comments (4)
Previously missed (3) — in code that hasn't changed since the last review.
cpp/src/arrow/array/array_primitive.h:149
- NumericArray::ToTensor() can pass a null Buffer into Tensor::Make when the values buffer is nullptr (e.g., an empty array built from ArrayData with buffers {nullptr, nullptr}). Tensor::Make rejects null data, so ToTensor() will fail for those empty arrays. Consider creating a non-null 0-byte Buffer when length()==0 and no buffer is present.
ARROW_ASSIGN_OR_RAISE(buffer, SliceBufferSafe(data_->buffers[1], boffset, blength));
}
return Tensor::Make(type(), std::move(buffer), {length()});
}
cpp/src/arrow/array/array_nested.cc:1046
- FixedSizeListArray::ToTensor() will fail for empty arrays if the leaf values buffer is nullptr (it forwards a null Buffer to Tensor::Make, which returns Invalid("Null data is supplied")). For length==0, it should be safe to use a non-null 0-byte Buffer instead so empty fixed-size-list tensors can still be created.
ARROW_ASSIGN_OR_RAISE(buffer, SliceBufferSafe(buf, boffset, blength));
}
return Tensor::Make(std::move(type), std::move(buffer), std::move(shape));
}
cpp/src/arrow/c/dlpack.cc:54
- GetDLDataType() is used for both Array and Tensor exports (ExportArrayImpl and ExportTensorImpl). The current error text suggests "try converting to a Tensor", which is confusing when the caller is already exporting a Tensor. Consider rewording the message so it remains accurate in both contexts (or move the hint to the array-only path).
return Status::TypeError(
"DataType is not compatible with DLPack spec: ", type.ToString(),
", try converting to a Tensor for multi dimensional data support");
}
cpp/src/arrow/c/dlpack_test.cc:174
- These assertions hard-code the exact DLPack type-compatibility error message. If GetDLDataType() is reworded to avoid implying callers should convert Tensors to tensors (since it is used by both array and tensor export paths), update both expected strings here to match the new wording.
ASSERT_RAISES_WITH_MESSAGE(TypeError,
"Type error: DataType is not compatible with DLPack spec: " +
array_null->type()->ToString() +
", try converting to a Tensor for multi"
" dimensional data support",
|
Thank you @AlenkaF, I reverted the unnecessary tensor changes. |
| remaining /= shape[i]; | ||
| strides->push_back(remaining); | ||
| // The outermost dimension is never a factor of the strides, so a shape whose total | ||
| // number of elements overflows can still have valid strides. |
There was a problem hiding this comment.
The strides would be valid, but computing an actual element address would overflow, so is it useful to allow this?
| if (internal::MultiplyWithOverflow(data_->offset, byte_width, &boffset) || | ||
| internal::MultiplyWithOverflow(length(), byte_width, &blength)) { | ||
| return Status::Invalid("Array byte size does not fit in an int64"); | ||
| } |
There was a problem hiding this comment.
That can't happen for a valid array, so we needn't check for this.
| /// Examples include NumericArray, FixedShapeTensorArray, nested FixedSizeListArray. | ||
| /// Nulls are ignored, leaving the output tensor with unspecified values where this | ||
| /// array has null entries. | ||
| virtual Result<std::shared_ptr<Tensor>> ToTensor() const; |
There was a problem hiding this comment.
API nit, but I think it would make more sense to expose Tensor facilities only in the corresponding headers, therefore have Tensor::FromArray rather than Array::ToTensor.
It would also mirror FixedShapeTensorArray::FromTensor.
| } | ||
|
|
||
| TEST_F(TestFixedSizeListArray, ToTensorNulls) { | ||
| // Nulls are ignored, leaving unspecified values in the output tensor. |
There was a problem hiding this comment.
Hmm... can we perhaps have an option to control that?
For example Tensor::FromArray(bool allow_nulls = false) or Array::ToTensor(bool allow_nulls = false)?
| if (internal::MultiplyWithOverflow(offset, int64_t{fsl->list_size()}, &offset) || | ||
| internal::AddWithOverflow(offset, data->offset, &offset) || | ||
| internal::MultiplyWithOverflow(length, int64_t{fsl->list_size()}, &length)) { |
There was a problem hiding this comment.
I think the overflow checks are not necessary here either. Overflow cannot happen on a valid array (because its data needs to fit in memory, therefore be smaller than INT64_MAX).
| } | ||
|
|
||
| TEST(TestPrimitiveArray, ToTensorNulls) { | ||
| // Nulls are ignored, leaving unspecified values in the output tensor. |
There was a problem hiding this comment.
Same comment as in array_list_test.cc.
Rationale for this change
Enable multidimensional DLPack support for Array via
to_tensor.What changes are included in this PR?
Array::ToTensorNumericArray::ToTensorfor 1D arraysFixedSizeListArray::ToTensorfor multidimensional arraysarr.to_tensor().__dlpack__()Note: Nulls are explicitly supported in
to_tensoras unspecified data. This was the current behaviour.Are these changes tested?
Yes
Are there any user-facing changes?
New public Array function.