diff --git a/cpp/src/arrow/sparse_tensor.cc b/cpp/src/arrow/sparse_tensor.cc index c23612e2530..0b67acd18c8 100644 --- a/cpp/src/arrow/sparse_tensor.cc +++ b/cpp/src/arrow/sparse_tensor.cc @@ -384,6 +384,21 @@ Result> SparseCSFIndex::Make( const std::vector>& indptr_data, const std::vector>& indices_data) { int64_t ndim = axis_order.size(); + + // Validate input dimensions before accessing buffers + if (ndim == 0) { + return Status::Invalid("axis_order cannot be empty"); + } + if (indices_shapes.size() != static_cast(ndim)) { + return Status::Invalid("indices_shapes must have the same size as axis_order"); + } + if (indptr_data.size() != static_cast(ndim - 1)) { + return Status::Invalid("indptr_data must have size equal to ndim - 1"); + } + if (indices_data.size() != static_cast(ndim)) { + return Status::Invalid("indices_data must have the same size as axis_order"); + } + std::vector> indptr(ndim - 1); std::vector> indices(ndim); diff --git a/cpp/src/arrow/sparse_tensor_test.cc b/cpp/src/arrow/sparse_tensor_test.cc index 434f4a1723c..ec440046eb8 100644 --- a/cpp/src/arrow/sparse_tensor_test.cc +++ b/cpp/src/arrow/sparse_tensor_test.cc @@ -1698,4 +1698,27 @@ TEST(TestSparseCSFMatrixForUInt64Index, Make) { ASSERT_RAISES(Invalid, SparseCSFTensor::Make(dense_tensor, uint64())); } +TEST(TestSparseCSFIndex, Make) { + std::vector axis_order = {0, 1, 2}; + std::vector indices_shapes = {5, 5, 5}; + std::vector> indptr_data(2); + std::vector> indices_data(3); + + // Invalid indices_shapes size + ASSERT_RAISES(Invalid, SparseCSFIndex::Make(int32(), int32(), {5}, axis_order, + indptr_data, indices_data)); + + // Empty axis order + ASSERT_RAISES(Invalid, SparseCSFIndex::Make(int32(), int32(), indices_shapes, {}, + indptr_data, indices_data)); + + // Invalid indptr_data size + ASSERT_RAISES(Invalid, SparseCSFIndex::Make(int32(), int32(), indices_shapes, + axis_order, {}, indices_data)); + + // Invalid indices_data size + ASSERT_RAISES(Invalid, SparseCSFIndex::Make(int32(), int32(), indices_shapes, + axis_order, indptr_data, {})); +} + } // namespace arrow