pitrou commented on code in PR #50929:
URL: https://github.com/apache/arrow/pull/50929#discussion_r3842710932
##########
cpp/src/arrow/tensor.cc:
##########
@@ -45,34 +45,36 @@ using internal::checked_cast;
namespace internal {
-Status ComputeRowMajorStrides(const FixedWidthType& type,
- const std::vector<int64_t>& shape,
- std::vector<int64_t>* strides) {
- const int byte_width = type.byte_width();
- const size_t ndim = shape.size();
-
- int64_t remaining = 0;
- if (!shape.empty() && shape.front() > 0) {
- remaining = byte_width;
- for (size_t i = 1; i < ndim; ++i) {
- if (internal::MultiplyWithOverflow(remaining, shape[i], &remaining)) {
- return Status::Invalid(
- "Row-major strides computed from shape would not fit in 64-bit
integer");
- }
- }
+Status ComputeRowMajorStrides(std::span<const int64_t> shape, int64_t
elem_size,
+ std::span<int64_t> strides) {
+ if (strides.size() != shape.size()) {
+ return Status::Invalid("strides must have the same length as shape");
}
- if (remaining == 0) {
- strides->assign(shape.size(), byte_width);
+ // An empty dimension makes the whole tensor empty, so any stride is as good.
+ if (std::find(shape.begin(), shape.end(), 0) != shape.end()) {
+ std::fill(strides.begin(), strides.end(), elem_size);
return Status::OK();
}
- strides->push_back(remaining);
- for (size_t i = 1; i < ndim; ++i) {
- 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.
Review Comment:
The strides would be valid, but computing an actual element address would
overflow, so is it useful to allow this?
##########
cpp/src/arrow/array/array_base.h:
##########
@@ -246,6 +248,15 @@ class ARROW_EXPORT Array {
/// \return const std::shared_ptr<ArrayStatistics>&
const std::shared_ptr<ArrayStatistics>& statistics() const { return
data_->statistics; }
+ /// \brief Create a Tensor from this Array
+ ///
+ /// When the data can reasonably be understood as a multidimensional numeric
Tensor,
+ /// return the data as such.
+ /// 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;
Review 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`.
##########
cpp/src/arrow/array/array_primitive.h:
##########
@@ -128,6 +131,23 @@ class NumericArray : public PrimitiveArray {
IteratorType end() const { return IteratorType(*this, length()); }
+ /// \brief Return a one dimensional Tensor.
+ Result<std::shared_ptr<Tensor>> ToTensor() const override {
+ // Could be non-templated
+ const int64_t byte_width = type()->byte_width();
+ std::shared_ptr<Buffer> buffer;
+ if (data_->buffers[1] != NULLPTR) {
+ int64_t boffset = 0;
+ int64_t blength = 0;
+ 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");
+ }
Review Comment:
That can't happen for a valid array, so we needn't check for this.
##########
cpp/src/arrow/array/array_nested.cc:
##########
@@ -1001,6 +1003,48 @@ Result<std::shared_ptr<Array>>
FixedSizeListArray::Flatten(
return FlattenListArray(*this, memory_pool);
}
+Result<std::shared_ptr<Tensor>> FixedSizeListArray::ToTensor() const {
+ const auto* data = this->data().get();
+ auto type = this->type();
+ int64_t offset = data->offset;
+ int64_t length = data->length;
+ std::vector<int64_t> shape{length};
+
+ // Iterate over nested fixed length container types.
+ // Each nested container increase the tensor dimension.
+ while (type->id() == Type::FIXED_SIZE_LIST) {
+ const auto* fsl = internal::checked_cast<const
FixedSizeListType*>(type.get());
+ type = fsl->value_type();
+ data = data->child_data.front().get();
+
+ 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)) {
Review 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).
##########
cpp/src/arrow/array/array_test.cc:
##########
@@ -1218,6 +1219,52 @@ TEST(TestPrimitiveArray, CtorNoValidityBitmap) {
ASSERT_EQ(arr.data()->null_count, 0);
}
+TEST(TestPrimitiveArray, ToTensor) {
+ const std::vector<int64_t> shape = {5};
+ const std::vector<int64_t> strides = {sizeof(int32_t)};
+
+ auto array = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5]");
+ ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+ ASSERT_OK(tensor->Validate());
+
+ EXPECT_EQ(int32(), tensor->type());
+ EXPECT_EQ(shape, tensor->shape());
+ EXPECT_EQ(strides, tensor->strides());
+ EXPECT_TRUE(tensor->is_contiguous());
+ EXPECT_TRUE(
+ TensorFromJSON(int32(), "[1, 2, 3, 4, 5]", shape,
strides)->Equals(*tensor));
+}
+
+TEST(TestPrimitiveArray, ToTensorSliced) {
+ const std::vector<int64_t> shape = {3};
+ const std::vector<int64_t> strides = {sizeof(int64_t)};
+
+ auto array = ArrayFromJSON(int64(), "[1, 2, 3, 4, 5]")->Slice(2);
+ ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+ ASSERT_OK(tensor->Validate());
+
+ EXPECT_EQ(shape, tensor->shape());
+ EXPECT_TRUE(TensorFromJSON(int64(), "[3, 4, 5]", shape,
strides)->Equals(*tensor));
+}
+
+TEST(TestPrimitiveArray, ToTensorNulls) {
+ // Nulls are ignored, leaving unspecified values in the output tensor.
Review Comment:
Same comment as in `array_list_test.cc`.
##########
cpp/src/arrow/array/array_list_test.cc:
##########
@@ -1821,4 +1822,106 @@ TEST_F(TestFixedSizeListArray, FlattenRecursively) {
*ArrayFromJSON(value_type_, "[0, 1, null, 3, 7, null, 2,
5]"));
}
+namespace {
+
+/// The flat values a tensor views: the innermost values of the nested fixed
size
+/// lists, windowed to what ``array`` covers.
+std::shared_ptr<Array> LeafValues(std::shared_ptr<Array> array) {
+ while (array->type_id() == Type::FIXED_SIZE_LIST) {
+ const auto& fsl = checked_cast<const FixedSizeListArray&>(*array);
+ array =
+ fsl.values()->Slice(fsl.value_offset(0), array->length() *
fsl.value_length());
+ }
+ return array;
+}
+
+template <typename T>
+void CheckToTensor(const std::shared_ptr<Array>& array, const
std::vector<int64_t>& shape,
+ std::initializer_list<T> values) {
+ const auto value_type = CTypeTraits<T>::type_singleton();
+ ASSERT_OK_AND_ASSIGN(
+ auto expected,
+ Tensor::Make(value_type, Buffer::Wrap(values.begin(), values.size()),
shape));
+
+ ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+ ASSERT_OK(tensor->Validate());
+
+ AssertTypeEqual(*value_type, *tensor->type());
+ ASSERT_EQ(shape, tensor->shape());
+ ASSERT_TRUE(tensor->is_row_major());
+ ASSERT_TRUE(tensor->Equals(*expected));
+
+ // The tensor shares the values buffer, it does not copy
+ const auto leaf = LeafValues(array);
+ ASSERT_EQ(leaf->data()->buffers[1]->data() + leaf->offset() * sizeof(T),
+ tensor->data()->data());
+}
+
+} // namespace
+
+TEST_F(TestFixedSizeListArray, ToTensor) {
+ auto array = ArrayFromJSON(fixed_size_list(int32(), 3),
+ "[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11,
12]]");
+ CheckToTensor<int32_t>(array, {4, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
12});
+
+ // Offset on the list array itself
+ CheckToTensor<int32_t>(array->Slice(2, 2), {2, 3}, {7, 8, 9, 10, 11, 12});
+
+ // Offset on the values array
+ auto values = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7, 8,
9]")->Slice(3);
+ ASSERT_OK_AND_ASSIGN(auto from_values,
FixedSizeListArray::FromArrays(values, 3));
+ CheckToTensor<int32_t>(from_values, {2, 3}, {4, 5, 6, 7, 8, 9});
+
+ // Offsets on both the list array and its values
+ CheckToTensor<int32_t>(from_values->Slice(1), {1, 3}, {7, 8, 9});
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorNested) {
+ auto array = ArrayFromJSON(fixed_size_list(fixed_size_list(float32(), 2),
3), R"([
+ [[1, 2], [3, 4], [5, 6]],
+ [[7, 8], [9, 10], [11, 12]],
+ [[13, 14], [15, 16], [17, 18]]
+ ])");
+ CheckToTensor<float>(array, {3, 3, 2},
+ {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
17, 18});
+
+ CheckToTensor<float>(array->Slice(1, 2), {2, 3, 2},
+ {7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18});
+
+ // 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.
+ auto values = ArrayFromJSON(float32(), "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
11, 12, 13]")
+ ->Slice(2, 12);
+ ASSERT_OK_AND_ASSIGN(auto inner, FixedSizeListArray::FromArrays(values, 2));
+ ASSERT_OK_AND_ASSIGN(auto outer, FixedSizeListArray::FromArrays(inner, 3));
+ CheckToTensor<float>(outer, {2, 3, 2}, {2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
13});
+ CheckToTensor<float>(outer->Slice(1), {1, 3, 2}, {8, 9, 10, 11, 12, 13});
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorNulls) {
+ // Nulls are ignored, leaving unspecified values in the output tensor.
Review 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)`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
