Skip to content

Commit 7548c3d

Browse files
GH-47552: [C++] Fix creating wrong object by FixedShapeTensorType::MakeArray() (#47533)
### Rationale for this change The builder should return the correct type. ### What changes are included in this PR? One liner ### Are these changes tested? Yes ### Are there any user-facing changes? No * GitHub Issue: #47552 Lead-authored-by: Fan Jiang <[email protected]> Co-authored-by: Fan Jiang <[email protected]> Signed-off-by: Sutou Kouhei <[email protected]>
1 parent 8145374 commit 7548c3d

File tree

2 files changed

+23
-1
lines changed

2 files changed

+23
-1
lines changed

cpp/src/arrow/extension/fixed_shape_tensor.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,7 @@ std::shared_ptr<Array> FixedShapeTensorType::MakeArray(
202202
DCHECK_EQ(data->type->id(), Type::EXTENSION);
203203
DCHECK_EQ("arrow.fixed_shape_tensor",
204204
internal::checked_cast<const ExtensionType&>(*data->type).extension_name());
205-
return std::make_shared<ExtensionArray>(data);
205+
return std::make_shared<FixedShapeTensorArray>(data);
206206
}
207207

208208
Result<std::shared_ptr<Tensor>> FixedShapeTensorType::MakeTensor(

cpp/src/arrow/extension/fixed_shape_tensor_test.cc

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,28 @@ TEST_F(TestExtensionType, CreateFromArray) {
152152
ASSERT_EQ(ext_arr->null_count(), 0);
153153
}
154154

155+
TEST_F(TestExtensionType, MakeArrayCanGetCorrectScalarType) {
156+
ASSERT_OK_AND_ASSIGN(auto tensor,
157+
Tensor::Make(value_type_, Buffer::Wrap(values_), shape_));
158+
159+
auto exact_ext_type = internal::checked_pointer_cast<FixedShapeTensorType>(ext_type_);
160+
ASSERT_OK_AND_ASSIGN(auto ext_arr, FixedShapeTensorArray::FromTensor(tensor));
161+
162+
auto data = ext_arr->data();
163+
auto array = internal::checked_pointer_cast<FixedShapeTensorArray>(
164+
exact_ext_type->MakeArray(data));
165+
ASSERT_EQ(array->length(), shape_[0]);
166+
ASSERT_EQ(array->null_count(), 0);
167+
168+
// Check that we can get the first element of the array
169+
ASSERT_OK_AND_ASSIGN(auto first_element, array->GetScalar(0));
170+
ASSERT_EQ(*(first_element->type),
171+
*(fixed_shape_tensor(value_type_, element_shape_, {0, 1})));
172+
173+
ASSERT_OK_AND_ASSIGN(auto tensor_from_array, array->ToTensor());
174+
ASSERT_TRUE(tensor->Equals(*tensor_from_array));
175+
}
176+
155177
void CheckSerializationRoundtrip(const std::shared_ptr<DataType>& ext_type) {
156178
auto fst_type = internal::checked_pointer_cast<FixedShapeTensorType>(ext_type);
157179
auto serialized = fst_type->Serialize();

0 commit comments

Comments
 (0)