jorisvandenbossche commented on code in PR #34797:
URL: https://github.com/apache/arrow/pull/34797#discussion_r1158138191
##########
cpp/src/arrow/extension/fixed_shape_tensor_test.cc:
##########
@@ -212,4 +222,204 @@ TEST_F(TestExtensionType, RoudtripBatch) {
CompareBatch(*batch, *read_batch2, /*compare_metadata=*/true);
}
+TEST_F(TestExtensionType, CreateFromTensor) {
+ std::vector<int64_t> column_major_strides = {8, 24, 72};
+ std::vector<int64_t> neither_major_strides = {96, 8, 32};
+
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ Tensor::Make(value_type_, Buffer::Wrap(values_),
shape_));
+
+ auto exact_ext_type =
internal::checked_pointer_cast<FixedShapeTensorType>(ext_type_);
+ ASSERT_OK_AND_ASSIGN(auto ext_arr,
FixedShapeTensorArray::FromTensor(tensor));
+
+ ASSERT_OK(ext_arr->ValidateFull());
+ ASSERT_TRUE(tensor->is_row_major());
+ ASSERT_EQ(tensor->strides(), tensor_strides_);
+ ASSERT_EQ(ext_arr->length(), shape_[0]);
+
+ auto ext_type_2 = internal::checked_pointer_cast<FixedShapeTensorType>(
+ fixed_shape_tensor(int64(), {3, 4}, {0, 1}));
+ ASSERT_OK_AND_ASSIGN(auto ext_arr_2,
FixedShapeTensorArray::FromTensor(tensor));
+
+ ASSERT_OK_AND_ASSIGN(
+ auto column_major_tensor,
+ Tensor::Make(value_type_, Buffer::Wrap(values_), shape_,
column_major_strides));
+ auto ext_type_3 = internal::checked_pointer_cast<FixedShapeTensorType>(
+ fixed_shape_tensor(int64(), {3, 4}, {0, 1}));
+ EXPECT_RAISES_WITH_MESSAGE_THAT(
+ Invalid,
+ testing::HasSubstr(
+ "Invalid: Only first-major tensors can be zero-copy converted to
arrays"),
+ FixedShapeTensorArray::FromTensor(column_major_tensor));
+ ASSERT_THAT(FixedShapeTensorArray::FromTensor(column_major_tensor),
+ Raises(StatusCode::Invalid));
+
+ auto neither_major_tensor = std::make_shared<Tensor>(value_type_,
Buffer::Wrap(values_),
+ shape_,
neither_major_strides);
+ auto ext_type_4 = internal::checked_pointer_cast<FixedShapeTensorType>(
+ fixed_shape_tensor(int64(), {3, 4}, {1, 0}));
+ ASSERT_OK_AND_ASSIGN(auto ext_arr_4,
+
FixedShapeTensorArray::FromTensor(neither_major_tensor));
+
+ auto ext_type_5 = internal::checked_pointer_cast<FixedShapeTensorType>(
+ fixed_shape_tensor(binary(), {1, 3}));
+ auto arr = ArrayFromJSON(binary(), R"(["abc", "def"])");
+
+ ASSERT_OK_AND_ASSIGN(auto fsla_arr,
+ FixedSizeListArray::FromArrays(arr,
fixed_size_list(binary(), 2)));
+ auto ext_arr_5 = std::reinterpret_pointer_cast<FixedShapeTensorArray>(
+ ExtensionType::WrapArray(ext_type_5, fsla_arr));
+ EXPECT_RAISES_WITH_MESSAGE_THAT(
+ Invalid, testing::HasSubstr("binary is not valid data type for a
tensor"),
+ ext_arr_5->ToTensor());
+
+ auto ext_type_6 = internal::checked_pointer_cast<FixedShapeTensorType>(
+ fixed_shape_tensor(int64(), {1, 2}));
+ auto arr_with_null = ArrayFromJSON(int64(), "[1, 0, null, null, 1, 2]");
+ ASSERT_OK_AND_ASSIGN(auto fsla_arr_6, FixedSizeListArray::FromArrays(
+ arr_with_null,
fixed_size_list(int64(), 2)));
+}
+
+void CheckFromTensorType(const std::shared_ptr<Tensor>& tensor,
+ std::shared_ptr<DataType> expected_ext_type) {
+ auto ext_type =
internal::checked_pointer_cast<FixedShapeTensorType>(expected_ext_type);
+ ASSERT_OK_AND_ASSIGN(auto ext_arr,
FixedShapeTensorArray::FromTensor(tensor));
+ auto generated_ext_type =
+ internal::checked_cast<const
FixedShapeTensorType*>(ext_arr->extension_type());
+
+ // Check that generated type is equal to the expected type
+ ASSERT_EQ(generated_ext_type->type_name(), ext_type->type_name());
+ ASSERT_EQ(generated_ext_type->shape(), ext_type->shape());
+ ASSERT_EQ(generated_ext_type->dim_names(), ext_type->dim_names());
+ ASSERT_EQ(generated_ext_type->permutation(), ext_type->permutation());
+
ASSERT_TRUE(generated_ext_type->storage_type()->Equals(*ext_type->storage_type()));
+ ASSERT_TRUE(generated_ext_type->Equals(ext_type));
+}
+
+TEST_F(TestExtensionType, TestFromTensorType) {
+ auto values = Buffer::Wrap(values_);
+ auto shapes =
+ std::vector<std::vector<int64_t>>{{3, 3, 4}, {3, 3, 4}, {3, 4, 3}, {3,
4, 3}};
+ auto strides = std::vector<std::vector<int64_t>>{
+ {96, 32, 8}, {96, 8, 24}, {96, 24, 8}, {96, 8, 32}};
+ auto cell_shapes = std::vector<std::vector<int64_t>>{{3, 4}, {3, 4}, {4, 3},
{4, 3}};
+ auto permutations = std::vector<std::vector<int64_t>>{{0, 1}, {1, 0}, {0,
1}, {1, 0}};
Review Comment:
So I assume this is the equivalent place compared to where I commented
before on the previous tests: for the second test case (shape {3, 3, 4},
cell_shape {3, 4} and permutation of {1, 0}), shouldn't the expected
fixed_shape_tensor type's shape be {4, 3}?
The tensor itself has shape {3, 3, 4} (so {3, 4} for the individual tensor
element), but has a permutation, so {4, 3} ?
--
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]
