Add partioned_matrix and tdb_partitioned_matrix unit tests & improve error handling (#381)

Author: jparismorgan

src/include/detail/linalg/matrix.h

@@ -461,8 +461,8 @@ constexpr auto SubMatrix(
 
 // TODO(paris): This only works on col-major matrices, fix for row-major.
 template <class Matrix>
-void debug_matrix(const Matrix& matrix, const std::string& msg = "") {
-  auto max_size = 10;
+void debug_matrix(
+    const Matrix& matrix, const std::string& msg = "", size_t max_size = 10) {
   auto rowsEnd = std::min(dimensions(matrix), static_cast<size_t>(max_size));
   auto colsEnd = std::min(num_vectors(matrix), static_cast<size_t>(max_size));
 

src/include/detail/linalg/matrix_with_ids.h

@@ -215,12 +215,13 @@ using ColMajorMatrixWithIds = MatrixWithIds<T, IdsType, stdx::layout_left, I>;
 // TODO(paris): This only works on col-major matrices, fix for row-major.
 template <class MatrixWithIds>
 void debug_matrix_with_ids(
-    const MatrixWithIds& matrix, const std::string& msg = "") {
-  auto max_size = 10;
+    const MatrixWithIds& matrix,
+    const std::string& msg = "",
+    size_t max_size = 10) {
   auto rowsEnd = std::min(dimensions(matrix), static_cast<size_t>(max_size));
   auto colsEnd = std::min(num_vectors(matrix), static_cast<size_t>(max_size));
 
-  debug_matrix(matrix);
+  debug_matrix(matrix, msg, max_size);
 
   std::cout << "# ids: [";
   auto end = std::min(matrix.num_ids(), static_cast<size_t>(max_size));

src/include/detail/linalg/partitioned_matrix.h

@@ -45,6 +45,7 @@
 
 #include <cstddef>
 #include "detail/linalg/matrix.h"
+#include "detail/linalg/vector.h"
 
 /**
  * @brief Partitioned matrix class.
@@ -154,6 +155,18 @@ class PartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       , part_index_(num_parts + 1)
       , num_vectors_{::num_vectors(training_set)}
       , num_parts_{num_parts} {
+    if (num_vectors_ == 0) {
+      throw std::invalid_argument("training_set cannot be empty.");
+    }
+    if (size(part_labels) != ::num_vectors(training_set)) {
+      throw std::invalid_argument(
+          "The number of part_labels must equal the number of vectors in the "
+          "training_set.");
+    }
+    if (num_parts <= 0) {
+      throw std::invalid_argument("num_parts should be greater than 0.");
+    }
+
     auto degrees = std::vector<size_t>(num_parts);
 
     for (size_t i = 0; i < ::num_vectors(training_set); ++i) {
@@ -203,10 +216,6 @@ class PartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
     return part_index_;
   }
 
-  auto& indices() {
-    return part_index_;
-  }
-
   virtual bool load() {
     return false;
   }
@@ -244,12 +253,13 @@ using ColMajorPartitionedMatrix = PartitionedMatrix<
 
 template <class PartitionedMatrix>
 void debug_partitioned_matrix(
-    const PartitionedMatrix& matrix, const std::string& msg = "") {
-  auto max_size = 10;
-  auto rowsEnd = std::min(dimension(matrix), static_cast<size_t>(max_size));
+    const PartitionedMatrix& matrix,
+    const std::string& msg = "",
+    size_t max_size = 10) {
+  auto rowsEnd = std::min(dimensions(matrix), static_cast<size_t>(max_size));
   auto colsEnd = std::min(num_vectors(matrix), static_cast<size_t>(max_size));
 
-  debug_matrix(matrix, msg);
+  debug_matrix(matrix, msg, max_size);
 
   std::cout << "# ids: [";
   auto end = std::min(matrix.ids().size(), static_cast<size_t>(max_size));

src/include/detail/linalg/tdb_partitioned_matrix.h

@@ -294,6 +294,14 @@ class tdbPartitionedMatrix
       , squashed_indices_(size(relevant_parts_) + 1)
       , first_resident_part_{0}
       , last_resident_part_{0} {
+    if (relevant_parts_.size() >= indices.size()) {
+      throw std::runtime_error(
+          "Invalid partitioning, relevant_parts_ size (" +
+          std::to_string(relevant_parts_.size()) +
+          ") must be less than indices size (" +
+          std::to_string(indices.size()) + ")");
+    }
+
     total_num_parts_ = size(relevant_parts_);
 
     scoped_timer _{tdb_func__ + " " + partitioned_vectors_uri_};
@@ -620,27 +628,7 @@ class tdbPartitionedMatrix
     num_loads_++;
     return true;
   }
-#if 0
-  auto& vectors() const {
-    return *this;
-  }
-
-  index_type num_resident_parts() const {
-    return last_resident_part_ - first_resident_part_;
-  }
 
-  index_type resident_part_offset() const {
-    return resident_part_offset_;
-  }
-
-  index_type col_offset() const {
-    return resident_col_offset_;
-  }
-
-  size_t num_loads() const {
-    return num_loads_;
-  }
-#endif
   /**
    * Destructor.  Closes arrays if they are open.
    */

src/include/detail/linalg/vector.h

@@ -161,26 +161,47 @@ class Vector : public std::span<T> {
 };
 
 template <feature_vector V>
-void debug_vector(const V& v, const std::string& msg = "") {
+void debug_vector(
+    const V& v, const std::string& msg = "", size_t max_size = 10) {
+  size_t end = std::min(max_size, dimensions(v));
   std::cout << msg << ": [";
-  for (size_t i = 0; i < dimensions(v); ++i) {
-    std::cout << v[i] << " ";
+  for (size_t i = 0; i < end; ++i) {
+    std::cout << v[i];
+    if (i != end - 1) {
+      std::cout << ", ";
+    }
+  }
+  if (dimensions(v) > max_size) {
+    std::cout << "...";
   }
   std::cout << "]\n";
 }
 
 template <std::ranges::forward_range V>
-void debug_vector(const V& v, const std::string& msg = "") {
+void debug_vector(
+    const V& v, const std::string& msg = "", size_t max_size = 10) {
+  size_t end = std::min(max_size, dimensions(v));
   std::cout << msg << ": [";
+  int idx = 0;
   for (auto&& i : v) {
-    std::cout << i << " ";
+    if (idx++ >= max_size) {
+      break;
+    }
+    std::cout << i;
+    if (i != end - 1) {
+      std::cout << ", ";
+    }
+  }
+  if (dimensions(v) > max_size) {
+    std::cout << "...";
   }
   std::cout << "]\n";
 }
 
 template <feature_vector V>
-void debug_matrix(const V& v, const std::string& msg = "") {
-  debug_vector(v, msg);
+void debug_matrix(
+    const V& v, const std::string& msg = "", size_t max_size = 10) {
+  debug_vector(v, msg, max_size);
 }
 
 #endif  // TILEDB_VECTOR_H

src/include/test/CMakeLists.txt

@@ -126,6 +126,8 @@ kmeans_add_test(unit_l2_distance)
 
 kmeans_add_test(unit_matrix)
 
+kmeans_add_test(unit_partitioned_matrix)
+
 kmeans_add_test(unit_matrix_with_ids)
 
 kmeans_add_test(unit_mdspan)
@@ -152,6 +154,8 @@ kmeans_add_test(unit_tdb_matrix)
 
 kmeans_add_test(unit_tdb_matrix_with_ids)
 
+kmeans_add_test(unit_tdb_partitioned_matrix)
+
 kmeans_add_test(unit_utils)
 
 kmeans_add_test(unit_vector)

src/include/test/unit_partitioned_matrix.cc

@@ -0,0 +1,144 @@
+/**
+ * @file   unit_partitioned_matrix.cc
+ *
+ * @section LICENSE
+ *
+ * The MIT License
+ *
+ * @copyright Copyright (c) 2024 TileDB, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ *
+ * @section DESCRIPTION
+ *
+ */
+
+#include <algorithm>
+#include <catch2/catch_all.hpp>
+#include <vector>
+#include "cpos.h"
+#include "detail/linalg/partitioned_matrix.h"
+#include "mdspan/mdspan.hpp"
+
+TEST_CASE("partitioned_matrix: test test", "[partitioned_matrix]") {
+  REQUIRE(true);
+}
+
+TEST_CASE("partitioned_matrix: sizes constructor", "[partitioned_matrix]") {
+  using feature_type = int;
+  using id_type = int;
+  using part_index_type = int;
+  size_t dimensions = 3;
+  size_t max_num_vectors = 5;
+  size_t max_num_partitions = 2;
+
+  auto partitioned_matrix =
+      ColMajorPartitionedMatrix<feature_type, id_type, part_index_type>(
+          dimensions, max_num_vectors, max_num_partitions);
+  CHECK(partitioned_matrix.num_vectors() == 0);
+  CHECK(partitioned_matrix.num_partitions() == 0);
+  CHECK(std::equal(
+      partitioned_matrix.ids().begin(),
+      partitioned_matrix.ids().end(),
+      std::vector<part_index_type>{0, 0, 0, 0, 0}.begin()));
+  CHECK(std::equal(
+      partitioned_matrix.indices().begin(),
+      partitioned_matrix.indices().end(),
+      std::vector<part_index_type>{0, 0, 0}.begin()));
+
+  CHECK(partitioned_matrix.load() == false);
+  CHECK(partitioned_matrix.num_vectors() == 0);
+  CHECK(partitioned_matrix.num_partitions() == 0);
+  CHECK(std::equal(
+      partitioned_matrix.ids().begin(),
+      partitioned_matrix.ids().end(),
+      std::vector<part_index_type>{0, 0, 0, 0, 0}.begin()));
+  CHECK(std::equal(
+      partitioned_matrix.indices().begin(),
+      partitioned_matrix.indices().end(),
+      std::vector<part_index_type>{0, 0, 0}.begin()));
+}
+
+TEST_CASE("partitioned_matrix: vectors constructor", "[partitioned_matrix]") {
+  using feature_type = float;
+  using id_type = float;
+  using part_index_type = float;
+
+  auto parts =
+      ColMajorMatrix<feature_type>{{1, 1, 1}, {2, 2, 2}, {3, 3, 3}, {4, 4, 4}};
+  std::vector<id_type> ids = {1, 2, 3, 4};
+  std::vector<part_index_type> part_index = {0, 1, 4};
+
+  auto partitioned_matrix =
+      ColMajorPartitionedMatrix<feature_type, id_type, part_index_type>(
+          parts, ids, part_index);
+
+  CHECK(partitioned_matrix.num_vectors() == 4);
+  CHECK(partitioned_matrix.num_partitions() == 2);
+  CHECK(std::equal(
+      partitioned_matrix.ids().begin(),
+      partitioned_matrix.ids().end(),
+      std::vector<part_index_type>{1, 2, 3, 4}.begin()));
+  CHECK(std::equal(
+      partitioned_matrix.indices().begin(),
+      partitioned_matrix.indices().end(),
+      std::vector<part_index_type>{0, 1, 4}.begin()));
+
+  CHECK(partitioned_matrix.load() == false);
+  CHECK(partitioned_matrix.num_vectors() == 4);
+  CHECK(partitioned_matrix.num_partitions() == 2);
+  CHECK(std::equal(
+      partitioned_matrix.ids().begin(),
+      partitioned_matrix.ids().end(),
+      std::vector<part_index_type>{1, 2, 3, 4}.begin()));
+  CHECK(std::equal(
+      partitioned_matrix.indices().begin(),
+      partitioned_matrix.indices().end(),
+      std::vector<part_index_type>{0, 1, 4}.begin()));
+}
+
+TEST_CASE("partitioned_matrix: training constructor", "[partitioned_matrix]") {
+  using feature_type = uint64_t;
+  using id_type = uint64_t;
+  using part_index_type = uint64_t;
+
+  auto training_set =
+      ColMajorMatrix<feature_type>{{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}};
+  std::vector<id_type> part_labels = {1, 0, 1, 0, 1};
+  size_t num_parts = 2;
+
+  auto partitioned_matrix =
+      ColMajorPartitionedMatrix<feature_type, id_type, part_index_type>(
+          training_set, part_labels, num_parts);
+  CHECK(partitioned_matrix.num_vectors() == _cpo::num_vectors(training_set));
+  CHECK(partitioned_matrix.num_partitions() == num_parts);
+  CHECK(std::equal(
+      partitioned_matrix.data(),
+      partitioned_matrix.data() + partitioned_matrix.num_vectors() *
+                                      _cpo::dimensions(partitioned_matrix),
+      std::vector<feature_type>{2, 2, 4, 4, 1, 1, 3, 3, 5, 5}.begin()));
+  CHECK(std::equal(
+      partitioned_matrix.ids().begin(),
+      partitioned_matrix.ids().end(),
+      std::vector<part_index_type>{1, 3, 0, 2, 4}.begin()));
+  CHECK(std::equal(
+      partitioned_matrix.indices().begin(),
+      partitioned_matrix.indices().end(),
+      std::vector<part_index_type>{0, 2, 5}.begin()));
+}