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 24 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")
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?
| /// 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.
There was a problem hiding this comment.
Sure, note however that in Python there is already array.to_tensor().
Also FixedShapeTensorArray::ToTensor is also an existing method.
There was a problem hiding this comment.
Ok, and I see we also have Table::ToTensor. So perhaps that ship has already sailed...
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
cpp/src/arrow/array/array_primitive.h:156
Tensor::Makerejects a nullBuffer("Null data is supplied"), but this implementation passes a nullbufferwhendata_->buffers[1]is NULL (which can be valid for empty fixed-width arrays). This makesToTensorunexpectedly fail for some valid empty arrays.
return Tensor::Make(type(), std::move(buffer), {length()});
cpp/src/arrow/array/array_nested.cc:1040
Tensor::Makerejects a nullBuffer, but this code leavesbufferas nullptr when the leaf values buffer is NULL (which can be valid when the computed leaf value count is 0). This can makeToTensorfail for some valid empty fixed-size-list arrays.
return Tensor::Make(std::move(type), std::move(buffer), std::move(shape));
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)
python/pyarrow/array.pxi:1856
- Docstring says allow_nulls defaults to True, but the function signature defaults to False. This mismatch will confuse users and docs generation.
allow_nulls: bool, default `True`
cpp/src/arrow/array/array_nested.cc:1035
- If the leaf values buffer is null for an empty FixedSizeListArray (possible for length==0 ArrayData), this passes a null data buffer to Tensor::Make, which fails validation. Consider creating an explicit 0-byte Buffer when the computed tensor size is 0.
std::shared_ptr<Buffer> buffer = nullptr;
if (const auto& buf = data->buffers[1]; buf != NULLPTR) {
const int64_t byte_width = type->byte_width();
// Buffer guarantees this fits into an int64_t.
const int64_t byte_offset = offset * byte_width;
cpp/src/arrow/array/array_primitive.h:152
- If a (valid) empty NumericArray has a null values buffer (buffers[1] == nullptr), this method passes a null data buffer into Tensor::Make, which fails validation even though the tensor has zero elements. Consider materializing an explicit 0-byte Buffer in that case.
std::shared_ptr<Buffer> buffer;
if (data_->buffers[1] != NULLPTR) {
// Array guarantees this will not overflow.
const int64_t byte_offset = data_->offset * byte_width;
const int64_t byte_length = length() * byte_width;
|
@pitrou I changed |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
python/pyarrow/array.pxi:1856
allow_nullsis documented as defaulting to True, but the signature defaults to False. This makes the docstring misleading and contradicts the behavior tested elsewhere (nulls rejected unless explicitly allowed).
Parameters
----------
allow_nulls: bool, default `True`
When true, nulls are ignored, leaving the output tensor with
python/pyarrow/array.pxi:4995
- This override of
FixedShapeTensorArray.to_tensor()shadowsArray.to_tensor(allow_nulls=...)but doesn't accept anallow_nullsargument. As a result,arr.to_tensor(allow_nulls=True)will raiseTypeErrorfor fixed-shape tensor arrays, breaking the new null-handling API and the added tests.
"""
return Array.to_tensor(self)
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
python/pyarrow/array.pxi:1856
- The docstring says
allow_nullsdefaults toTrue, but the Python signature defaults toFalse(and the C++ default is alsofalse). This is user-facing API documentation, so it should match the actual default behavior.
allow_nulls: bool, default `True`
python/pyarrow/array.pxi:4995
- This delegation uses
Array.to_tensor(self)without exposing the newallow_nullsparameter. ForFixedShapeTensorArrayinstances,arr.to_tensor(allow_nulls=True)will raise a PythonTypeErrorbecause this override’s signature isto_tensor(self)only.
return Array.to_tensor(self)
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 (2)
cpp/src/arrow/c/dlpack.cc:54
- Typo/wording: "multi dimensional" should be hyphenated as "multi-dimensional" in this user-facing error message (and update test expectations accordingly).
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/array/array_nested.cc:1023
offset = offset * list_size + ...andlength = length * list_sizecan trigger signed overflow (UB) on large arrays beforeSliceBufferSafehas a chance to validate bounds. Use overflow-checked arithmetic to keep this safe even on malformed/unvalidated inputs.
// Overflow cannot happen on a valid array (its data needs to fit in memory,
// therefore be smaller than INT64_MAX)
offset = offset * fsl->list_size() + data->offset;
length = length * fsl->list_size();
shape.push_back(fsl->list_size());
pitrou
left a comment
pitrou
left a comment
There was a problem hiding this comment.
A bunch of nits and minor suggestions, but LGTM in general!
|
|
||
| namespace arrow { | ||
|
|
||
| class Tensor; |
There was a problem hiding this comment.
Can include arrow/type_fwd.h to avoid adding declarations like this.
| // Offsets accumulate at every level: the innermost values are offset by 2, the | ||
| // middle lists by 3 * 2 and the outer lists by 1 * 3 * 2. |
There was a problem hiding this comment.
The middle lists do not seem sliced below?
| // Nulls are ignored, leaving unspecified values in the output tensor. | ||
| ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor(/* allow_nulls= */ true)); | ||
| ASSERT_OK(tensor->Validate()); | ||
| ASSERT_EQ(std::vector<int64_t>({3, 2}), tensor->shape()); |
There was a problem hiding this comment.
We can probably check the 3 non-null tensor values?
| NotImplemented, | ||
| ArrayFromJSON(fixed_size_list(utf8(), 1), R"([["a"], ["b"]])")->ToTensor()); | ||
| ASSERT_RAISES( | ||
| Invalid, | ||
| ArrayFromJSON(fixed_size_list(boolean(), 2), "[[true, false]]")->ToTensor()); |
There was a problem hiding this comment.
Why do we get NotImplemented and Invalid? Ideally this should return TypeError.
| ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor(/* allow_nulls= */ true)); | ||
| ASSERT_OK(tensor->Validate()); | ||
|
|
||
| EXPECT_EQ(shape, tensor->shape()); |
There was a problem hiding this comment.
We can also check the non-null values here?
|
|
||
| namespace arrow { | ||
|
|
||
| class Array; |
There was a problem hiding this comment.
Same here: can include arrow/type_fwd.h instead.
| array = array.copy() | ||
| return array | ||
|
|
||
| def to_tensor(self, allow_nulls=False): |
There was a problem hiding this comment.
We probably want to make the boolean flag allow_nulls keyword-only to discourage bad practices:
| def to_tensor(self, allow_nulls=False): | |
| def to_tensor(self, *, allow_nulls=False): |
| return self.to_tensor().to_numpy() | ||
|
|
||
| def to_tensor(self): | ||
| def to_tensor(self, allow_nulls=False): |
There was a problem hiding this comment.
Same here: make it keyword-only?
| def to_tensor(self, allow_nulls=False): | |
| def to_tensor(self, *, allow_nulls=False): |
| with nogil: | ||
| ctensor = ext_array.ToTensor() | ||
| return pyarrow_wrap_tensor(GetResultValue(ctensor)) | ||
| return Array.to_tensor(self, allow_nulls=allow_nulls) |
There was a problem hiding this comment.
Can we just inherit this method or do you keep it for the more concrete docstring?
| # A Tensor sharing an Array buffer is immutable, so it can only be exported | ||
| # through the versioned DLPack protocol. | ||
| assert not tensor.is_mutable | ||
| result = np.from_dlpack(DLPackForwarder(tensor, max_version=(1, 0))) |
There was a problem hiding this comment.
Is it possible to also call np.from_dlpack(arr) or is that not possible yet?
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.