diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index 4a61c9115774b..15e6f8d731fcd 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -335,13 +335,16 @@ const Result> FixedShapeTensorArray::ToTensor() const { // To convert an array of n dimensional tensors to a n+1 dimensional tensor we // interpret the array's length as the first dimension the new tensor. - const auto ext_arr = + const auto storage_array = internal::checked_pointer_cast(this->storage()); + ARROW_ASSIGN_OR_RAISE(const auto flattened_storage_array, storage_array->Flatten()); + const auto ext_type = internal::checked_pointer_cast(this->type()); - ARROW_RETURN_IF(!is_fixed_width(*ext_arr->value_type()), - Status::Invalid(ext_arr->value_type()->ToString(), - " is not valid data type for a tensor")); + const auto value_type = ext_type->value_type(); + ARROW_RETURN_IF( + !is_fixed_width(*value_type), + Status::Invalid(value_type->ToString(), " is not valid data type for a tensor")); std::vector permutation = ext_type->permutation(); if (permutation.empty()) { @@ -361,18 +364,21 @@ const Result> FixedShapeTensorArray::ToTensor() const { } std::vector shape = ext_type->shape(); + auto cell_size = std::accumulate(shape.begin(), shape.end(), static_cast(1), + std::multiplies<>()); shape.insert(shape.begin(), 1, this->length()); internal::Permute(permutation, &shape); std::vector tensor_strides; - const auto value_type = - internal::checked_pointer_cast(ext_arr->value_type()); + const auto fw_value_type = internal::checked_pointer_cast(value_type); ARROW_RETURN_NOT_OK( - ComputeStrides(*value_type.get(), shape, permutation, &tensor_strides)); - ARROW_ASSIGN_OR_RAISE(const auto flattened_array, ext_arr->Flatten()); + ComputeStrides(*fw_value_type.get(), shape, permutation, &tensor_strides)); - return Tensor::Make(ext_arr->value_type(), flattened_array->data()->buffers[1], shape, - tensor_strides, dim_names); + ARROW_ASSIGN_OR_RAISE( + const auto buffer, + SliceBufferSafe(flattened_storage_array->data()->buffers[1], + this->offset() * cell_size * value_type->byte_width())); + return Tensor::Make(value_type, buffer, shape, tensor_strides, dim_names); } Result> FixedShapeTensorType::Make( diff --git a/python/pyarrow/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py index 5892ffc9c3696..0e0c4fff47a0f 100644 --- a/python/pyarrow/tests/test_extension_type.py +++ b/python/pyarrow/tests/test_extension_type.py @@ -1334,8 +1334,7 @@ def test_tensor_class_methods(value_type): expected = np.array([[[7, 8, 9], [10, 11, 12]]], dtype=value_type) result = arr[1:].to_numpy_ndarray() - # TODO: offset of sliced pa.array is not correctly handled - # np.testing.assert_array_equal(result, expected) + np.testing.assert_array_equal(result, expected) values = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] flat_arr = np.array(values[0], dtype=value_type) @@ -1413,8 +1412,7 @@ def test_tensor_array_from_numpy(value_type): arr = np.array([[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]], dtype=value_type) expected = arr[1:] result = pa.FixedShapeTensorArray.from_numpy_ndarray(arr)[1:].to_numpy_ndarray() - # TODO: offset of sliced pa.array is not correctly handled - # np.testing.assert_array_equal(result, expected) + np.testing.assert_array_equal(result, expected) @pytest.mark.parametrize("tensor_type", (