GH-50816: [C++][IPC] Improve sparse tensor index validation

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<flatbuf::Buffer> indptr, indices;
 
   for (const std::shared_ptr<arrow::Tensor>& 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<arrow::Tensor>& 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 (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(
+          "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<int>(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();

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);
@@ -589,6 +593,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,

cpp/src/arrow/sparse_tensor.cc

@@ -221,13 +221,15 @@ Result<std::shared_ptr<SparseCOOIndex>> SparseCOOIndex::Make(
     const std::shared_ptr<Tensor>& coords, bool is_canonical) {
   RETURN_NOT_OK(
       CheckSparseCOOIndexValidity(coords->type(), coords->shape(), coords->strides()));
+  RETURN_NOT_OK(coords->Validate());
   return std::make_shared<SparseCOOIndex>(coords, is_canonical);
 }
 
 Result<std::shared_ptr<SparseCOOIndex>> SparseCOOIndex::Make(
     const std::shared_ptr<Tensor>& 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<SparseCOOIndex>(coords, is_canonical);
 }
@@ -239,10 +241,10 @@ Result<std::shared_ptr<SparseCOOIndex>> SparseCOOIndex::Make(
     bool is_canonical) {
   RETURN_NOT_OK(
       CheckSparseCOOIndexValidity(indices_type, indices_shape, indices_strides));
-  return std::make_shared<SparseCOOIndex>(
-      std::make_shared<Tensor>(indices_type, indices_data, indices_shape,
-                               indices_strides),
-      is_canonical);
+  auto coords = std::make_shared<Tensor>(indices_type, indices_data, indices_shape,
+                                         indices_strides);
+  RETURN_NOT_OK(coords->Validate());
+  return std::make_shared<SparseCOOIndex>(coords, is_canonical);
 }
 
 Result<std::shared_ptr<SparseCOOIndex>> SparseCOOIndex::Make(
@@ -253,6 +255,7 @@ Result<std::shared_ptr<SparseCOOIndex>> SparseCOOIndex::Make(
       CheckSparseCOOIndexValidity(indices_type, indices_shape, indices_strides));
   auto coords = std::make_shared<Tensor>(indices_type, indices_data, indices_shape,
                                          indices_strides);
+  RETURN_NOT_OK(coords->Validate());
   auto is_canonical = DetectSparseCOOIndexCanonicality(coords);
   return std::make_shared<SparseCOOIndex>(coords, is_canonical);
 }
@@ -293,6 +296,29 @@ SparseCOOIndex::SparseCOOIndex(const std::shared_ptr<Tensor>& coords, bool is_ca
 
 std::string SparseCOOIndex::ToString() const { return std::string("SparseCOOIndex"); }
 
+Status SparseCOOIndex::ValidateShape(const std::vector<int64_t>& shape) const {
+  RETURN_NOT_OK(SparseIndex::ValidateShape(shape));
+  RETURN_NOT_OK(coords_->Validate());
+
+  if (coords_->ndim() != 2 || static_cast<size_t>(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));
+
+  std::vector<int64_t> 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) {
+      if (index[dim] < 0 || index[dim] >= shape[dim]) {
+        return Status::Invalid("SparseCOOIndex index is out of bounds for dimension ",
+                               dim);
+      }
+    }
+  }
+  return Status::OK();
+}
+
 // ----------------------------------------------------------------------
 // SparseCSXIndex
 
@@ -337,6 +363,61 @@ Status ValidateSparseCSXIndex(const std::shared_ptr<DataType>& indptr_type,
   return Status::OK();
 }
 
+Status ValidateSparseCSXIndexContents(SparseMatrixCompressedAxis compressed_axis,
+                                      const std::shared_ptr<Tensor>& indptr,
+                                      const std::shared_ptr<Tensor>& indices,
+                                      const std::vector<int64_t>& 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<DataType>& indptr_type,
                                  const std::shared_ptr<DataType>& indices_type,
                                  const std::vector<int64_t>& indptr_shape,
@@ -357,7 +438,7 @@ inline Status CheckSparseCSFIndexValidity(const std::shared_ptr<DataType>& indpt
                                           const std::shared_ptr<DataType>& indices_type,
                                           const int64_t num_indptrs,
                                           const int64_t num_indices,
-                                          const int64_t axis_order_size) {
+                                          const std::vector<int64_t>& 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<DataType>& 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<int64_t>(axis_order.size()) != num_indices) {
     return Status::Invalid(
         "Length of indices must be equal to number of dimensions for SparseCSFIndex.");
   }
+  std::vector<bool> 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<std::shared_ptr<SparseCSFIndex>> SparseCSFIndex::Make(
                                           std::vector<int64_t>({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<std::shared_ptr<Tensor>>& 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<int64_t>& 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<int64_t>(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<int64_t>(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<std::shared_ptr<Tensor>> 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<int64_t>(
+      {non_zero_length(), static_cast<int64_t>(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(