Add tdbPreLoadMatrixWithIds (#284)

Author: jparismorgan

src/include/detail/linalg/tdb_matrix_with_ids.h

@@ -241,6 +241,63 @@ class tdbBlockedMatrixWithIds
   }
 };  // tdbBlockedMatrixWithIds
 
+template <
+    class T,
+    class IdsType = uint64_t,
+    class LayoutPolicy = stdx::layout_right,
+    class I = size_t>
+class tdbPreLoadMatrixWithIds
+    : public tdbBlockedMatrixWithIds<T, IdsType, LayoutPolicy, I> {
+  using Base = tdbBlockedMatrixWithIds<T, IdsType, LayoutPolicy, I>;
+
+ public:
+  /**
+   * @brief Construct a new tdbBlockedMatrixWithIds object, limited to
+   * `upper_bound` vectors. In this case, the `Matrix` is column-major, so the
+   * number of vectors is the number of columns.
+   *
+   * @param ctx The TileDB context to use.
+   * @param uri URI of the TileDB array to read.
+   * @param upper_bound The maximum number of vectors to read.
+   */
+  tdbPreLoadMatrixWithIds(
+      const tiledb::Context& ctx,
+      const std::string& uri,
+      const std::string& ids_uri,
+      size_t upper_bound = 0,
+      uint64_t timestamp = 0)
+      : tdbPreLoadMatrixWithIds(
+            ctx,
+            uri,
+            ids_uri,
+            std::nullopt,
+            std::nullopt,
+            upper_bound,
+            timestamp) {
+  }
+
+  tdbPreLoadMatrixWithIds(
+      const tiledb::Context& ctx,
+      const std::string& uri,
+      const std::string& ids_uri,
+      std::optional<size_t> num_array_rows,
+      std::optional<size_t> num_array_cols,
+      size_t upper_bound = 0,
+      uint64_t timestamp = 0)
+      : Base(
+            ctx,
+            uri,
+            ids_uri,
+            0,
+            num_array_rows,
+            0,
+            num_array_cols,
+            upper_bound,
+            timestamp) {
+    Base::load();
+  }
+};
+
 /**
  * Convenience class for row-major blocked matrices.
  */
@@ -279,4 +336,18 @@ template <
     class I = size_t>
 using tdbMatrixWithIds = tdbBlockedMatrixWithIds<T, IdsType, LayoutPolicy, I>;
 
+/**
+ * Convenience class for row-major matrices.
+ */
+template <class T, class IdsType = uint64_t, class I = size_t>
+using tdbRowMajorPreLoadMatrixWithIds =
+    tdbPreLoadMatrixWithIds<T, IdsType, stdx::layout_right, I>;
+
+/**
+ * Convenience class for column-major matrices.
+ */
+template <class T, class IdsType = uint64_t, class I = size_t>
+using tdbColMajorPreLoadMatrixWithIds =
+    tdbPreLoadMatrixWithIds<T, IdsType, stdx::layout_left, I>;
+
 #endif  // TDB_MATRIX_WITH_IDS_H

src/include/test/unit_tdb_matrix_with_ids.cc

@@ -308,3 +308,67 @@ TEST_CASE("tdb_matrix_with_ids: empty matrix", "[tdb_matrix_with_ids]") {
     CHECK(Y.ids().size() == 1000);
   }
 }
+
+TEMPLATE_TEST_CASE(
+    "tdb_matrix_with_ids: preload", "[tdb_matrix_with_ids]", float, uint8_t) {
+  tiledb::Context ctx;
+  std::string tmp_matrix_uri =
+      (std::filesystem::temp_directory_path() / "tmp_tdb_matrix").string();
+  std::string tmp_ids_uri =
+      (std::filesystem::temp_directory_path() / "tmp_tdb_ids_matrix").string();
+  int offset = 13;
+  size_t Mrows = 200;
+  size_t Ncols = 500;
+
+  tiledb::VFS vfs(ctx);
+  if (vfs.is_dir(tmp_matrix_uri)) {
+    vfs.remove_dir(tmp_matrix_uri);
+  }
+
+  auto X = ColMajorMatrixWithIds<TestType, TestType>(Mrows, Ncols);
+  fill_and_write_matrix(
+      ctx, X, tmp_matrix_uri, tmp_ids_uri, Mrows, Ncols, offset);
+  CHECK(X.ids()[0] == offset + 0);
+  CHECK(X.ids()[1] == offset + 1);
+  CHECK(X.ids()[10] == offset + 10);
+
+  auto Y = tdbPreLoadMatrixWithIds<TestType, TestType, stdx::layout_left>(
+      ctx, tmp_matrix_uri, tmp_ids_uri);
+  CHECK(num_vectors(Y) == num_vectors(X));
+  CHECK(dimension(Y) == dimension(X));
+  CHECK(
+      std::equal(X.data(), X.data() + dimension(X) * num_vectors(X), Y.data()));
+  for (size_t i = 0; i < 5; ++i) {
+    for (size_t j = 0; j < 5; ++j) {
+      CHECK(X(i, j) == Y(i, j));
+    }
+  }
+
+  CHECK(size(Y.ids()) == Y.num_ids());
+  CHECK(size(X.ids()) == X.num_ids());
+  CHECK(X.num_ids() == Y.num_ids());
+  CHECK(std::equal(X.ids().begin(), X.ids().end(), Y.ids().begin()));
+  for (size_t i = 0; i < X.num_ids(); ++i) {
+    CHECK(X.ids()[i] == Y.ids()[i]);
+  }
+
+  auto Z = ColMajorMatrixWithIds<TestType, TestType>(0, 0);
+  Z = std::move(Y);
+
+  CHECK(num_vectors(Z) == num_vectors(X));
+  CHECK(dimension(Z) == dimension(X));
+  CHECK(
+      std::equal(X.data(), X.data() + dimension(X) * num_vectors(X), Z.data()));
+  for (size_t i = 0; i < 5; ++i) {
+    for (size_t j = 0; j < 5; ++j) {
+      CHECK(X(i, j) == Z(i, j));
+    }
+  }
+
+  CHECK(size(Z.ids()) == Z.num_ids());
+  CHECK(X.num_ids() == Z.num_ids());
+  CHECK(std::equal(X.ids().begin(), X.ids().end(), Z.ids().begin()));
+  for (size_t i = 0; i < X.num_ids(); ++i) {
+    CHECK(X.ids()[i] == Z.ids()[i]);
+  }
+}