From 36bfb2255646bcaeacb8e6f0c55f0c550ba04ef1 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 5 Aug 2026 20:16:29 +0200 Subject: [PATCH 1/4] [C++][IPC] Validate sparse tensor indices before conversion --- cpp/src/arrow/ipc/reader.cc | 22 +++- cpp/src/arrow/ipc/tensor_test.cc | 12 ++ cpp/src/arrow/sparse_tensor.cc | 185 +++++++++++++++++++++++++++- cpp/src/arrow/sparse_tensor.h | 36 ++---- cpp/src/arrow/sparse_tensor_test.cc | 53 ++++++++ 5 files changed, 274 insertions(+), 34 deletions(-) diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc index 47ea70e43fac..de1de7718daa 100644 --- a/cpp/src/arrow/ipc/reader.cc +++ b/cpp/src/arrow/ipc/reader.cc @@ -2438,14 +2438,24 @@ Result> ReadSparseCSFIndex( RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata( sparse_index, &axis_order, &indices_size, &indptr_type, &indices_type)); for (int i = 0; i < static_cast(indptr_buffers->size()); ++i) { - ARROW_ASSIGN_OR_RAISE(indptr_data[i], file->ReadAt(indptr_buffers->Get(i)->offset(), - indptr_buffers->Get(i)->length(), - /*allow_short_read=*/false)); + const auto byte_length = MultiplyWithOverflow( + {indptr_buffers->Get(i)->length(), indptr_type->byte_width()}); + if (!byte_length.has_value()) { + return Status::Invalid("SparseCSFIndex indptr buffer size overflows"); + } + ARROW_ASSIGN_OR_RAISE(indptr_data[i], + file->ReadAt(indptr_buffers->Get(i)->offset(), + byte_length.value(), /*allow_short_read=*/false)); } for (int i = 0; i < static_cast(indices_buffers->size()); ++i) { - ARROW_ASSIGN_OR_RAISE(indices_data[i], file->ReadAt(indices_buffers->Get(i)->offset(), - indices_buffers->Get(i)->length(), - /*allow_short_read=*/false)); + const auto byte_length = MultiplyWithOverflow( + {indices_buffers->Get(i)->length(), indices_type->byte_width()}); + if (!byte_length.has_value()) { + return Status::Invalid("SparseCSFIndex indices buffer size overflows"); + } + ARROW_ASSIGN_OR_RAISE(indices_data[i], + file->ReadAt(indices_buffers->Get(i)->offset(), + byte_length.value(), /*allow_short_read=*/false)); } return SparseCSFIndex::Make(indptr_type, indices_type, indices_size, axis_order, diff --git a/cpp/src/arrow/ipc/tensor_test.cc b/cpp/src/arrow/ipc/tensor_test.cc index a9243e779932..770ec98c40d9 100644 --- a/cpp/src/arrow/ipc/tensor_test.cc +++ b/cpp/src/arrow/ipc/tensor_test.cc @@ -589,6 +589,18 @@ TEST(TestSparseCSFIndex, RejectInconsistentBufferCounts) { ASSERT_RAISES(Invalid, ReadSparseTensor(*message)); } +TEST(TestSparseCSFIndex, RejectInvalidAxisOrder) { + // MakeCSFSparseTensorMessage initializes axisOrder to all zeroes, which is not a + // permutation for a two-dimensional tensor. + ASSERT_OK_AND_ASSIGN(auto message, MakeCSFSparseTensorMessage(/*shape=*/{4, 4}, + /*num_indptr_buffers=*/1, + /*num_indices_buffers=*/2, + /*axis_order_size=*/2)); + ASSERT_RAISES(Invalid, ReadSparseTensor(*message)); + ASSERT_RAISES(Invalid, + internal::ReadSparseTensorPayload(MakeSparseTensorPayload(message, 4))); +} + TEST(TestSparseCSFIndex, RejectInconsistentPayloadBufferCounts) { ASSERT_OK_AND_ASSIGN(auto message, MakeCSFSparseTensorMessage(/*shape=*/{4}, /*num_indptr_buffers=*/0, diff --git a/cpp/src/arrow/sparse_tensor.cc b/cpp/src/arrow/sparse_tensor.cc index c23612e25306..6cf494f60291 100644 --- a/cpp/src/arrow/sparse_tensor.cc +++ b/cpp/src/arrow/sparse_tensor.cc @@ -293,6 +293,32 @@ SparseCOOIndex::SparseCOOIndex(const std::shared_ptr& coords, bool is_ca std::string SparseCOOIndex::ToString() const { return std::string("SparseCOOIndex"); } +Status SparseCOOIndex::ValidateShape(const std::vector& shape) const { + RETURN_NOT_OK(SparseIndex::ValidateShape(shape)); + RETURN_NOT_OK(coords_->Validate()); + + if (coords_->ndim() != 2 || static_cast(coords_->shape()[1]) != shape.size()) { + return Status::Invalid( + "shape length is inconsistent with the coords matrix in COO index"); + } + RETURN_NOT_OK(internal::CheckSparseIndexMaximumValue(coords_->type(), shape)); + + const int index_elsize = coords_->type()->byte_width(); + const uint8_t* coords_data = coords_->raw_data(); + for (int64_t i = 0; i < coords_->shape()[0]; ++i) { + for (int64_t dim = 0; dim < coords_->shape()[1]; ++dim) { + const int64_t index = + internal::SparseTensorConverterMixin::GetIndexValue(coords_data, index_elsize); + if (index < 0 || index >= shape[dim]) { + return Status::Invalid("SparseCOOIndex index is out of bounds for dimension ", + dim); + } + coords_data += index_elsize; + } + } + return Status::OK(); +} + // ---------------------------------------------------------------------- // SparseCSXIndex @@ -337,6 +363,61 @@ Status ValidateSparseCSXIndex(const std::shared_ptr& indptr_type, return Status::OK(); } +Status ValidateSparseCSXIndexContents(SparseMatrixCompressedAxis compressed_axis, + const std::shared_ptr& indptr, + const std::shared_ptr& indices, + const std::vector& shape, + const char* type_name) { + if (shape.size() != 2) { + return Status::Invalid("Invalid shape length for a sparse matrix"); + } + RETURN_NOT_OK(ValidateSparseCSXIndex(indptr->type(), indices->type(), indptr->shape(), + indices->shape(), type_name)); + RETURN_NOT_OK(indptr->Validate()); + RETURN_NOT_OK(indices->Validate()); + RETURN_NOT_OK(CheckSparseIndexMaximumValue(indices->type(), shape)); + + ARROW_ASSIGN_OR_RAISE(const int64_t expected_indptr_length, + ComputeSparseCSXIndptrLength(compressed_axis, shape)); + if (indptr->shape()[0] != expected_indptr_length) { + return Status::Invalid("shape is inconsistent with the ", type_name); + } + + const int64_t non_zero_length = indices->shape()[0]; + const int indptr_elsize = indptr->type()->byte_width(); + const uint8_t* indptr_data = indptr->raw_data(); + int64_t previous = + SparseTensorConverterMixin::GetIndexValue(indptr_data, indptr_elsize); + if (previous != 0) { + return Status::Invalid(type_name, " indptr must start at 0"); + } + for (int64_t i = 1; i < expected_indptr_length; ++i) { + const int64_t current = SparseTensorConverterMixin::GetIndexValue( + indptr_data + i * indptr_elsize, indptr_elsize); + if (current < previous || current > non_zero_length) { + return Status::Invalid(type_name, + " indptr values must be non-decreasing and not exceed " + "the indices length"); + } + previous = current; + } + if (previous != non_zero_length) { + return Status::Invalid(type_name, " indptr must end at the indices length"); + } + + const int64_t minor_axis = compressed_axis == SparseMatrixCompressedAxis::ROW ? 1 : 0; + const int indices_elsize = indices->type()->byte_width(); + const uint8_t* indices_data = indices->raw_data(); + for (int64_t i = 0; i < non_zero_length; ++i) { + const int64_t index = SparseTensorConverterMixin::GetIndexValue( + indices_data + i * indices_elsize, indices_elsize); + if (index < 0 || index >= shape[minor_axis]) { + return Status::Invalid(type_name, " index is out of bounds"); + } + } + return Status::OK(); +} + void CheckSparseCSXIndexValidity(const std::shared_ptr& indptr_type, const std::shared_ptr& indices_type, const std::vector& indptr_shape, @@ -357,7 +438,7 @@ inline Status CheckSparseCSFIndexValidity(const std::shared_ptr& indpt const std::shared_ptr& indices_type, const int64_t num_indptrs, const int64_t num_indices, - const int64_t axis_order_size) { + const std::vector& axis_order) { if (!is_integer(indptr_type->id())) { return Status::TypeError("Type of SparseCSFIndex indptr must be integer"); } @@ -368,10 +449,18 @@ inline Status CheckSparseCSFIndexValidity(const std::shared_ptr& indpt return Status::Invalid( "Length of indices must be equal to length of indptrs + 1 for SparseCSFIndex."); } - if (axis_order_size != num_indices) { + if (static_cast(axis_order.size()) != num_indices) { return Status::Invalid( "Length of indices must be equal to number of dimensions for SparseCSFIndex."); } + std::vector seen(num_indices, false); + for (const int64_t axis : axis_order) { + if (axis < 0 || axis >= num_indices || seen[axis]) { + return Status::Invalid( + "SparseCSFIndex axis_order must be a permutation of the dimensions"); + } + seen[axis] = true; + } return Status::OK(); } @@ -395,7 +484,7 @@ Result> SparseCSFIndex::Make( std::vector({indices_shapes[i]})); RETURN_NOT_OK(CheckSparseCSFIndexValidity(indptr_type, indices_type, indptr.size(), - indices.size(), axis_order.size())); + indices.size(), axis_order)); for (auto tensor : indptr) { RETURN_NOT_OK(internal::CheckSparseIndexMaximumValue(indptr_type, tensor->shape())); @@ -415,11 +504,75 @@ SparseCSFIndex::SparseCSFIndex(const std::vector>& indpt : SparseIndexBase(), indptr_(indptr), indices_(indices), axis_order_(axis_order) { ARROW_CHECK_OK(CheckSparseCSFIndexValidity(indptr_.front()->type(), indices_.front()->type(), indptr_.size(), - indices_.size(), axis_order_.size())); + indices_.size(), axis_order_)); } std::string SparseCSFIndex::ToString() const { return std::string("SparseCSFIndex"); } +Status SparseCSFIndex::ValidateShape(const std::vector& shape) const { + RETURN_NOT_OK(SparseIndex::ValidateShape(shape)); + RETURN_NOT_OK(CheckSparseCSFIndexValidity(indptr_.front()->type(), + indices_.front()->type(), indptr_.size(), + indices_.size(), axis_order_)); + if (shape.size() != axis_order_.size()) { + return Status::Invalid("shape length is inconsistent with the SparseCSFIndex"); + } + RETURN_NOT_OK(internal::CheckSparseIndexMaximumValue(indices_.front()->type(), shape)); + + for (int64_t dim = 0; dim < static_cast(indices_.size()); ++dim) { + const auto& cur_indices = indices_[dim]; + if (cur_indices->ndim() != 1) { + return Status::Invalid("SparseCSFIndex indices must be vectors"); + } + RETURN_NOT_OK(cur_indices->Validate()); + + const int indices_elsize = cur_indices->type()->byte_width(); + const uint8_t* indices_data = cur_indices->raw_data(); + for (int64_t i = 0; i < cur_indices->shape()[0]; ++i) { + const int64_t index = internal::SparseTensorConverterMixin::GetIndexValue( + indices_data + i * indices_elsize, indices_elsize); + if (index < 0 || index >= shape[axis_order_[dim]]) { + return Status::Invalid("SparseCSFIndex index is out of bounds for dimension ", + axis_order_[dim]); + } + } + + if (dim == static_cast(indptr_.size())) { + continue; + } + const auto& cur_indptr = indptr_[dim]; + if (cur_indptr->ndim() != 1 || + cur_indptr->shape()[0] != cur_indices->shape()[0] + 1) { + return Status::Invalid( + "SparseCSFIndex indptr length is inconsistent with its indices"); + } + RETURN_NOT_OK(cur_indptr->Validate()); + + const int indptr_elsize = cur_indptr->type()->byte_width(); + const uint8_t* indptr_data = cur_indptr->raw_data(); + int64_t previous = + internal::SparseTensorConverterMixin::GetIndexValue(indptr_data, indptr_elsize); + if (previous != 0) { + return Status::Invalid("SparseCSFIndex indptr must start at 0"); + } + const int64_t next_indices_length = indices_[dim + 1]->shape()[0]; + for (int64_t i = 1; i < cur_indptr->shape()[0]; ++i) { + const int64_t current = internal::SparseTensorConverterMixin::GetIndexValue( + indptr_data + i * indptr_elsize, indptr_elsize); + if (current < previous || current > next_indices_length) { + return Status::Invalid( + "SparseCSFIndex indptr values must be non-decreasing and not exceed " + "the next indices length"); + } + previous = current; + } + if (previous != next_indices_length) { + return Status::Invalid("SparseCSFIndex indptr must end at the next indices length"); + } + } + return Status::OK(); +} + bool SparseCSFIndex::Equals(const SparseCSFIndex& other) const { auto eq = [](const auto& a, const auto& b) { return a->Equals(*b); }; // TODO: remove the use of std::equal when we no longer have partial C++20 support with @@ -473,6 +626,30 @@ bool SparseTensor::Equals(const SparseTensor& other, const EqualOptions& opts) c } Result> SparseTensor::ToTensor(MemoryPool* pool) const { + if (!sparse_index_) { + return Status::Invalid("Sparse tensor has no sparse index"); + } + if (!data_) { + return Status::Invalid("Sparse tensor has no values buffer"); + } + RETURN_NOT_OK(sparse_index_->ValidateShape(shape_)); + + const auto values_size = internal::MultiplyWithOverflow( + {non_zero_length(), static_cast(type_->byte_width())}); + if (!values_size.has_value() || values_size.value() > data_->size()) { + return Status::Invalid("Sparse tensor values buffer is too small"); + } + + int64_t dense_size = 1; + for (const int64_t dim_size : shape_) { + if (internal::MultiplyWithOverflow(dense_size, dim_size, &dense_size)) { + return Status::Invalid("Sparse tensor size exceeds the maximum supported size"); + } + } + if (internal::MultiplyWithOverflow(dense_size, type_->byte_width(), &dense_size)) { + return Status::Invalid("Sparse tensor size exceeds the maximum supported size"); + } + switch (format_id()) { case SparseTensorFormat::COO: return MakeTensorFromSparseCOOTensor( diff --git a/cpp/src/arrow/sparse_tensor.h b/cpp/src/arrow/sparse_tensor.h index a41024b74160..6f14c3a2b8de 100644 --- a/cpp/src/arrow/sparse_tensor.h +++ b/cpp/src/arrow/sparse_tensor.h @@ -177,16 +177,7 @@ class ARROW_EXPORT SparseCOOIndex : public internal::SparseIndexBaseEquals(*other.indices()); } - inline Status ValidateShape(const std::vector& shape) const override { - ARROW_RETURN_NOT_OK(SparseIndex::ValidateShape(shape)); - - if (static_cast(coords_->shape()[1]) == shape.size()) { - return Status::OK(); - } - - return Status::Invalid( - "shape length is inconsistent with the coords matrix in COO index"); - } + Status ValidateShape(const std::vector& shape) const override; protected: std::shared_ptr coords_; @@ -214,6 +205,13 @@ Status ValidateSparseCSXIndex(const std::shared_ptr& indptr_type, const std::vector& indices_shape, const char* type_name); +ARROW_EXPORT +Status ValidateSparseCSXIndexContents(SparseMatrixCompressedAxis compressed_axis, + const std::shared_ptr& indptr, + const std::shared_ptr& indices, + const std::vector& shape, + const char* type_name); + ARROW_EXPORT void CheckSparseCSXIndexValidity(const std::shared_ptr& indptr_type, const std::shared_ptr& indices_type, @@ -304,20 +302,8 @@ class SparseCSXIndex : public SparseIndexBase { inline Status ValidateShape(const std::vector& shape) const override { ARROW_RETURN_NOT_OK(SparseIndex::ValidateShape(shape)); - - if (shape.size() < 2) { - return Status::Invalid("shape length is too short"); - } - - if (shape.size() > 2) { - return Status::Invalid("shape length is too long"); - } - - if (indptr_->shape()[0] == shape[static_cast(kCompressedAxis)] + 1) { - return Status::OK(); - } - - return Status::Invalid("shape length is inconsistent with the ", ToString()); + return ValidateSparseCSXIndexContents(kCompressedAxis, indptr_, indices_, shape, + SparseIndexType::kTypeName); } protected: @@ -448,6 +434,8 @@ class ARROW_EXPORT SparseCSFIndex : public internal::SparseIndexBase& shape) const override; + protected: std::vector> indptr_; std::vector> indices_; diff --git a/cpp/src/arrow/sparse_tensor_test.cc b/cpp/src/arrow/sparse_tensor_test.cc index 434f4a1723c7..6ce1e763cb25 100644 --- a/cpp/src/arrow/sparse_tensor_test.cc +++ b/cpp/src/arrow/sparse_tensor_test.cc @@ -1050,6 +1050,24 @@ TEST(TestSparseCSRMatrixForUInt64Index, Make) { ASSERT_RAISES(Invalid, SparseCSRMatrix::Make(dense_tensor, uint64())); } +TEST(TestSparseCSRMatrix, RejectOutOfBoundsIndex) { + std::vector indptr_values = {0, 1, 1}; + std::vector indices_values = {3}; + ASSERT_OK_AND_ASSIGN( + auto index, SparseCSRIndex::Make(int64(), /*indptr_shape=*/{3}, + /*indices_shape=*/{1}, Buffer::Wrap(indptr_values), + Buffer::Wrap(indices_values))); + std::vector sparse_values = {1}; + auto sparse_tensor = std::make_shared( + index, int64(), Buffer::Wrap(sparse_values), std::vector{2, 3}, + std::vector{}); + + ASSERT_RAISES(Invalid, sparse_tensor->ToTensor()); + ASSERT_RAISES(Invalid, + SparseCSRMatrix::Make(index, int64(), Buffer::Wrap(sparse_values), + /*shape=*/{2, 3}, /*dim_names=*/{})); +} + template class TestSparseCSCMatrixBase : public TestSparseTensorBase { public: @@ -1698,4 +1716,39 @@ TEST(TestSparseCSFMatrixForUInt64Index, Make) { ASSERT_RAISES(Invalid, SparseCSFTensor::Make(dense_tensor, uint64())); } +TEST(TestSparseCSFIndex, RejectInvalidAxisOrder) { + std::vector indptr_values = {0, 1}; + std::vector first_indices_values = {0}; + std::vector second_indices_values = {0}; + std::vector> indptr = {Buffer::Wrap(indptr_values)}; + std::vector> indices = {Buffer::Wrap(first_indices_values), + Buffer::Wrap(second_indices_values)}; + + ASSERT_RAISES(Invalid, SparseCSFIndex::Make(int64(), /*indices_shapes=*/{1, 1}, + /*axis_order=*/{0, 2}, indptr, indices)); + ASSERT_RAISES(Invalid, SparseCSFIndex::Make(int64(), /*indices_shapes=*/{1, 1}, + /*axis_order=*/{0, 0}, indptr, indices)); +} + +TEST(TestSparseCSFTensor, RejectOutOfBoundsIndex) { + std::vector indptr_values = {0, 1}; + std::vector first_indices_values = {0}; + std::vector second_indices_values = {3}; + std::vector> indptr = {Buffer::Wrap(indptr_values)}; + std::vector> indices = {Buffer::Wrap(first_indices_values), + Buffer::Wrap(second_indices_values)}; + ASSERT_OK_AND_ASSIGN(auto index, + SparseCSFIndex::Make(int64(), /*indices_shapes=*/{1, 1}, + /*axis_order=*/{0, 1}, indptr, indices)); + std::vector sparse_values = {1}; + auto sparse_tensor = std::make_shared( + index, int64(), Buffer::Wrap(sparse_values), std::vector{2, 3}, + std::vector{}); + + ASSERT_RAISES(Invalid, sparse_tensor->ToTensor()); + ASSERT_RAISES(Invalid, + SparseCSFTensor::Make(index, int64(), Buffer::Wrap(sparse_values), + /*shape=*/{2, 3}, /*dim_names=*/{})); +} + } // namespace arrow From 40a7803c40c9cc86ba652810d7dd0f760f3b3cef Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 5 Aug 2026 20:39:43 +0200 Subject: [PATCH 2/4] [C++][IPC] Preserve sparse index layout and buffer semantics --- cpp/src/arrow/ipc/metadata_internal.cc | 41 ++++++++++++++++++-------- cpp/src/arrow/ipc/reader.cc | 22 ++++---------- cpp/src/arrow/ipc/tensor_test.cc | 12 +++++--- cpp/src/arrow/sparse_tensor.cc | 20 ++++++------- cpp/src/arrow/sparse_tensor_test.cc | 15 ++++++++++ cpp/src/arrow/tensor/coo_converter.cc | 8 +++-- 6 files changed, 73 insertions(+), 45 deletions(-) diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc index 94a956916162..35c5a909c3eb 100644 --- a/cpp/src/arrow/ipc/metadata_internal.cc +++ b/cpp/src/arrow/ipc/metadata_internal.cc @@ -1163,22 +1163,19 @@ Status MakeSparseTensorIndexCSF(FBB& fbb, const SparseCSFIndex& sparse_index, auto indices_type_offset = flatbuf::CreateInt(fbb, indices_value_type.bit_width(), indices_value_type.is_signed()); - const int64_t indptr_elem_size = indptr_value_type.byte_width(); - const int64_t indices_elem_size = indices_value_type.byte_width(); - int64_t offset = 0; std::vector indptr, indices; for (const std::shared_ptr& tensor : sparse_index.indptr()) { - const int64_t size = tensor->data()->size() / indptr_elem_size; - const int64_t padded_size = PaddedLength(tensor->data()->size(), kArrowIpcAlignment); + const int64_t size = tensor->data()->size(); + const int64_t padded_size = PaddedLength(size, kArrowIpcAlignment); indptr.push_back({offset, size}); offset += padded_size; } for (const std::shared_ptr& tensor : sparse_index.indices()) { - const int64_t size = tensor->data()->size() / indices_elem_size; - const int64_t padded_size = PaddedLength(tensor->data()->size(), kArrowIpcAlignment); + const int64_t size = tensor->data()->size(); + const int64_t padded_size = PaddedLength(size, kArrowIpcAlignment); indices.push_back({offset, size}); offset += padded_size; @@ -1523,19 +1520,39 @@ Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_ind RETURN_NOT_OK(IntFromFlatbuffer(sparse_index->indicesType(), indices_type)); auto* fb_axis_order = sparse_index->axisOrder(); + auto* fb_indptr_buffers = sparse_index->indptrBuffers(); auto* fb_indices_buffers = sparse_index->indicesBuffers(); // ValidateSparseCSFIndexMetadata already checks this, keep this check defensively. - if (fb_axis_order == nullptr || fb_indices_buffers == nullptr || - fb_axis_order->size() != fb_indices_buffers->size()) { + if (fb_axis_order == nullptr || fb_indptr_buffers == nullptr || + fb_indices_buffers == nullptr || + fb_axis_order->size() != fb_indices_buffers->size() || + fb_indptr_buffers->size() + 1 != fb_indices_buffers->size()) { return Status::Invalid( - "Inconsistent CSF sparse index: axisOrder and indicesBuffers have different " - "lengths"); + "Inconsistent CSF sparse index: indptrBuffers, indicesBuffers and axisOrder " + "have different lengths"); } + const int64_t indptr_byte_width = (*indptr_type)->byte_width(); + for (size_t i = 0; i < fb_indptr_buffers->size(); ++i) { + const int64_t byte_length = fb_indptr_buffers->Get(i)->length(); + if (byte_length < 0 || byte_length % indptr_byte_width != 0) { + return Status::Invalid( + "SparseCSFIndex indptr buffer size must be a non-negative multiple of " + "the index element size"); + } + } + + const int64_t indices_byte_width = (*indices_type)->byte_width(); const int ndim = static_cast(fb_axis_order->size()); for (int i = 0; i < ndim; ++i) { + const int64_t byte_length = fb_indices_buffers->Get(i)->length(); + if (byte_length < 0 || byte_length % indices_byte_width != 0) { + return Status::Invalid( + "SparseCSFIndex indices buffer size must be a non-negative multiple of " + "the index element size"); + } axis_order->push_back(fb_axis_order->Get(i)); - indices_size->push_back(fb_indices_buffers->Get(i)->length()); + indices_size->push_back(byte_length / indices_byte_width); } return Status::OK(); diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc index de1de7718daa..47ea70e43fac 100644 --- a/cpp/src/arrow/ipc/reader.cc +++ b/cpp/src/arrow/ipc/reader.cc @@ -2438,24 +2438,14 @@ Result> ReadSparseCSFIndex( RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata( sparse_index, &axis_order, &indices_size, &indptr_type, &indices_type)); for (int i = 0; i < static_cast(indptr_buffers->size()); ++i) { - const auto byte_length = MultiplyWithOverflow( - {indptr_buffers->Get(i)->length(), indptr_type->byte_width()}); - if (!byte_length.has_value()) { - return Status::Invalid("SparseCSFIndex indptr buffer size overflows"); - } - ARROW_ASSIGN_OR_RAISE(indptr_data[i], - file->ReadAt(indptr_buffers->Get(i)->offset(), - byte_length.value(), /*allow_short_read=*/false)); + ARROW_ASSIGN_OR_RAISE(indptr_data[i], file->ReadAt(indptr_buffers->Get(i)->offset(), + indptr_buffers->Get(i)->length(), + /*allow_short_read=*/false)); } for (int i = 0; i < static_cast(indices_buffers->size()); ++i) { - const auto byte_length = MultiplyWithOverflow( - {indices_buffers->Get(i)->length(), indices_type->byte_width()}); - if (!byte_length.has_value()) { - return Status::Invalid("SparseCSFIndex indices buffer size overflows"); - } - ARROW_ASSIGN_OR_RAISE(indices_data[i], - file->ReadAt(indices_buffers->Get(i)->offset(), - byte_length.value(), /*allow_short_read=*/false)); + ARROW_ASSIGN_OR_RAISE(indices_data[i], file->ReadAt(indices_buffers->Get(i)->offset(), + indices_buffers->Get(i)->length(), + /*allow_short_read=*/false)); } return SparseCSFIndex::Make(indptr_type, indices_type, indices_size, axis_order, diff --git a/cpp/src/arrow/ipc/tensor_test.cc b/cpp/src/arrow/ipc/tensor_test.cc index 770ec98c40d9..f6b496369fd7 100644 --- a/cpp/src/arrow/ipc/tensor_test.cc +++ b/cpp/src/arrow/ipc/tensor_test.cc @@ -268,12 +268,16 @@ class TestSparseTensorRoundTrip : public BaseTensorTest { int64_t out_indptr_length = 0; int64_t out_indices_length = 0; for (int i = 0; i < ndim - 1; ++i) { - out_indptr_length += bit_util::RoundUpToMultipleOf8( - index_elem_size * resulted_sparse_index.indptr()[i]->size()); + const auto& indptr = resulted_sparse_index.indptr()[i]; + ASSERT_EQ(index_elem_size * indptr->size(), indptr->data()->size()); + out_indptr_length += + bit_util::RoundUpToMultipleOf8(index_elem_size * indptr->size()); } for (int i = 0; i < ndim; ++i) { - out_indices_length += bit_util::RoundUpToMultipleOf8( - index_elem_size * resulted_sparse_index.indices()[i]->size()); + const auto& indices = resulted_sparse_index.indices()[i]; + ASSERT_EQ(index_elem_size * indices->size(), indices->data()->size()); + out_indices_length += + bit_util::RoundUpToMultipleOf8(index_elem_size * indices->size()); } ASSERT_EQ(out_indptr_length, indptr_length); diff --git a/cpp/src/arrow/sparse_tensor.cc b/cpp/src/arrow/sparse_tensor.cc index 6cf494f60291..0a0b71bc89d1 100644 --- a/cpp/src/arrow/sparse_tensor.cc +++ b/cpp/src/arrow/sparse_tensor.cc @@ -221,6 +221,7 @@ Result> SparseCOOIndex::Make( const std::shared_ptr& coords, bool is_canonical) { RETURN_NOT_OK( CheckSparseCOOIndexValidity(coords->type(), coords->shape(), coords->strides())); + RETURN_NOT_OK(coords->Validate()); return std::make_shared(coords, is_canonical); } @@ -228,6 +229,7 @@ Result> SparseCOOIndex::Make( const std::shared_ptr& coords) { RETURN_NOT_OK( CheckSparseCOOIndexValidity(coords->type(), coords->shape(), coords->strides())); + RETURN_NOT_OK(coords->Validate()); auto is_canonical = DetectSparseCOOIndexCanonicality(coords); return std::make_shared(coords, is_canonical); } @@ -239,10 +241,10 @@ Result> SparseCOOIndex::Make( bool is_canonical) { RETURN_NOT_OK( CheckSparseCOOIndexValidity(indices_type, indices_shape, indices_strides)); - return std::make_shared( - std::make_shared(indices_type, indices_data, indices_shape, - indices_strides), - is_canonical); + auto coords = std::make_shared(indices_type, indices_data, indices_shape, + indices_strides); + RETURN_NOT_OK(coords->Validate()); + return std::make_shared(coords, is_canonical); } Result> SparseCOOIndex::Make( @@ -253,6 +255,7 @@ Result> SparseCOOIndex::Make( CheckSparseCOOIndexValidity(indices_type, indices_shape, indices_strides)); auto coords = std::make_shared(indices_type, indices_data, indices_shape, indices_strides); + RETURN_NOT_OK(coords->Validate()); auto is_canonical = DetectSparseCOOIndexCanonicality(coords); return std::make_shared(coords, is_canonical); } @@ -303,17 +306,14 @@ Status SparseCOOIndex::ValidateShape(const std::vector& shape) const { } RETURN_NOT_OK(internal::CheckSparseIndexMaximumValue(coords_->type(), shape)); - const int index_elsize = coords_->type()->byte_width(); - const uint8_t* coords_data = coords_->raw_data(); + std::vector index; for (int64_t i = 0; i < coords_->shape()[0]; ++i) { + GetCOOIndexTensorRow(coords_, i, &index); for (int64_t dim = 0; dim < coords_->shape()[1]; ++dim) { - const int64_t index = - internal::SparseTensorConverterMixin::GetIndexValue(coords_data, index_elsize); - if (index < 0 || index >= shape[dim]) { + if (index[dim] < 0 || index[dim] >= shape[dim]) { return Status::Invalid("SparseCOOIndex index is out of bounds for dimension ", dim); } - coords_data += index_elsize; } } return Status::OK(); diff --git a/cpp/src/arrow/sparse_tensor_test.cc b/cpp/src/arrow/sparse_tensor_test.cc index 6ce1e763cb25..635d85e279b1 100644 --- a/cpp/src/arrow/sparse_tensor_test.cc +++ b/cpp/src/arrow/sparse_tensor_test.cc @@ -153,6 +153,14 @@ TEST(TestSparseCOOIndex, MakeEmptyIndex) { ASSERT_TRUE(si->is_canonical()); } +TEST(TestSparseCOOIndex, RejectShortIndicesBuffer) { + auto data = Buffer::FromString(""); + ASSERT_RAISES(Invalid, + SparseCOOIndex::Make( + int32(), /*indices_shape=*/{1, 2}, + /*indices_strides=*/{2 * sizeof(int32_t), sizeof(int32_t)}, data)); +} + TEST(TestSparseCSRIndex, Make) { std::vector indptr_values = {0, 2, 4, 6, 8, 10, 12}; std::vector indices_values = {0, 2, 1, 3, 0, 2, 1, 3, 0, 2, 1, 3}; @@ -658,6 +666,13 @@ TYPED_TEST_P(TestSparseCOOTensorForIndexValueType, CreationWithColumnMajorIndex) ASSERT_EQ("baz", st->dim_name(2)); ASSERT_TRUE(st->Equals(*this->sparse_tensor_from_dense_)); + + ASSERT_OK_AND_ASSIGN(auto dense_tensor, st->ToTensor()); + std::vector dense_values = {1, 0, 2, 0, 0, 3, 0, 4, 5, 0, 6, 0, + 0, 11, 0, 12, 13, 0, 14, 0, 0, 15, 0, 16}; + Tensor expected(int64(), Buffer::Wrap(dense_values), this->shape_, {}, + this->dim_names_); + ASSERT_TRUE(dense_tensor->Equals(expected)); } TYPED_TEST_P(TestSparseCOOTensorForIndexValueType, diff --git a/cpp/src/arrow/tensor/coo_converter.cc b/cpp/src/arrow/tensor/coo_converter.cc index 7e29b668f53e..b835a5377a8e 100644 --- a/cpp/src/arrow/tensor/coo_converter.cc +++ b/cpp/src/arrow/tensor/coo_converter.cc @@ -294,6 +294,7 @@ Result> MakeTensorFromSparseCOOTensor( checked_cast(*sparse_tensor->sparse_index()); const auto& coords = sparse_index.indices(); const auto* coords_data = coords->raw_data(); + const auto& coords_strides = coords->strides(); const int index_elsize = coords->type()->byte_width(); @@ -314,10 +315,11 @@ Result> MakeTensorFromSparseCOOTensor( int64_t offset = 0; for (int j = 0; j < ndim; ++j) { - auto index = static_cast( - SparseTensorConverterMixin::GetIndexValue(coords_data, index_elsize)); + const auto* index_data = + coords_data + i * coords_strides[0] + j * coords_strides[1]; + const auto index = static_cast( + SparseTensorConverterMixin::GetIndexValue(index_data, index_elsize)); offset += index * strides[j]; - coords_data += index_elsize; } std::copy_n(raw_data, value_elsize, values + offset); From e647bfadf0c8dd783c14b40ebf361f81151c8c71 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 5 Aug 2026 20:58:47 +0200 Subject: [PATCH 3/4] [C++][IPC] Fix sparse tensor CI failures --- cpp/src/arrow/ipc/metadata_internal.cc | 2 +- cpp/src/arrow/sparse_tensor_test.cc | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc index 35c5a909c3eb..bd156b8cf719 100644 --- a/cpp/src/arrow/ipc/metadata_internal.cc +++ b/cpp/src/arrow/ipc/metadata_internal.cc @@ -1533,7 +1533,7 @@ Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_ind } const int64_t indptr_byte_width = (*indptr_type)->byte_width(); - for (size_t i = 0; i < fb_indptr_buffers->size(); ++i) { + for (flatbuffers::uoffset_t i = 0; i < fb_indptr_buffers->size(); ++i) { const int64_t byte_length = fb_indptr_buffers->Get(i)->length(); if (byte_length < 0 || byte_length % indptr_byte_width != 0) { return Status::Invalid( diff --git a/cpp/src/arrow/sparse_tensor_test.cc b/cpp/src/arrow/sparse_tensor_test.cc index 635d85e279b1..35ad16e60804 100644 --- a/cpp/src/arrow/sparse_tensor_test.cc +++ b/cpp/src/arrow/sparse_tensor_test.cc @@ -1065,7 +1065,7 @@ TEST(TestSparseCSRMatrixForUInt64Index, Make) { ASSERT_RAISES(Invalid, SparseCSRMatrix::Make(dense_tensor, uint64())); } -TEST(TestSparseCSRMatrix, RejectOutOfBoundsIndex) { +TEST(TestSparseCSRMatrixValidation, RejectOutOfBoundsIndex) { std::vector indptr_values = {0, 1, 1}; std::vector indices_values = {3}; ASSERT_OK_AND_ASSIGN( @@ -1745,7 +1745,7 @@ TEST(TestSparseCSFIndex, RejectInvalidAxisOrder) { /*axis_order=*/{0, 0}, indptr, indices)); } -TEST(TestSparseCSFTensor, RejectOutOfBoundsIndex) { +TEST(TestSparseCSFTensorValidation, RejectOutOfBoundsIndex) { std::vector indptr_values = {0, 1}; std::vector first_indices_values = {0}; std::vector second_indices_values = {3}; From 473cec3bc972842b7acaa424b750d5e00ea6a6d4 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 5 Aug 2026 21:44:07 +0200 Subject: [PATCH 4/4] [C++][IPC] Disambiguate sparse index test overload --- cpp/src/arrow/sparse_tensor_test.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/sparse_tensor_test.cc b/cpp/src/arrow/sparse_tensor_test.cc index 35ad16e60804..cd6d209eddab 100644 --- a/cpp/src/arrow/sparse_tensor_test.cc +++ b/cpp/src/arrow/sparse_tensor_test.cc @@ -1068,10 +1068,12 @@ TEST(TestSparseCSRMatrixForUInt64Index, Make) { TEST(TestSparseCSRMatrixValidation, RejectOutOfBoundsIndex) { std::vector indptr_values = {0, 1, 1}; std::vector indices_values = {3}; + const std::vector indptr_shape = {3}; + const std::vector indices_shape = {1}; ASSERT_OK_AND_ASSIGN( - auto index, SparseCSRIndex::Make(int64(), /*indptr_shape=*/{3}, - /*indices_shape=*/{1}, Buffer::Wrap(indptr_values), - Buffer::Wrap(indices_values))); + auto index, + SparseCSRIndex::Make(int64(), indptr_shape, indices_shape, + Buffer::Wrap(indptr_values), Buffer::Wrap(indices_values))); std::vector sparse_values = {1}; auto sparse_tensor = std::make_shared( index, int64(), Buffer::Wrap(sparse_values), std::vector{2, 3},