rok commented on code in PR #34797:
URL: https://github.com/apache/arrow/pull/34797#discussion_r1160127242
##########
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:
Changed `ToTensor/FromTensor` to permute original shape to cell_shape and
back.
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
