Support writing C++ type-erased indexes with a storage_version (#326)

Author: jparismorgan

apis/python/src/tiledb/vector_search/type_erased_module.cc

@@ -283,11 +283,13 @@ void init_type_erased_module(py::module_& m) {
           "write_index",
           [](IndexVamana& index,
              const tiledb::Context& ctx,
-             const std::string& group_uri) {
-            index.write_index(ctx, group_uri);
+             const std::string& group_uri,
+             const std::string& storage_version) {
+            index.write_index(ctx, group_uri, storage_version);
           },
           py::arg("ctx"),
-          py::arg("group_uri"))
+          py::arg("group_uri"),
+          py::arg("storage_version") = "")
       .def("feature_type_string", &IndexVamana::feature_type_string)
       .def("id_type_string", &IndexVamana::id_type_string)
       .def(

apis/python/src/tiledb/vector_search/vamana_index.py

@@ -97,7 +97,6 @@ def query_internal(
         return np.array(distances, copy=False), np.array(ids, copy=False)
 
 
-# TODO(paris): Pass more arguments to C++, i.e. storage_version.
 def create(
     uri: str,
     dimensions: int,
@@ -120,5 +119,5 @@ def create(
     )
     index.train(empty_vector)
     index.add(empty_vector)
-    index.write_index(ctx, uri)
+    index.write_index(ctx, uri, storage_version)
     return VamanaIndex(uri=uri, config=config, memory_budget=1000000)

src/include/api/ivf_flat_index.h

@@ -427,12 +427,14 @@ class IndexIVFFlat {
   }
 
   void write_index(
-      const tiledb::Context& ctx, const std::string& group_uri) const {
+      const tiledb::Context& ctx,
+      const std::string& group_uri,
+      const std::string& storage_version = "") const {
     if (!index_) {
       throw std::runtime_error(
           "Cannot write_index() because there is no index.");
     }
-    index_->write_index(ctx, group_uri);
+    index_->write_index(ctx, group_uri, storage_version);
   }
 
   constexpr auto dimension() const {
@@ -503,7 +505,9 @@ class IndexIVFFlat {
     virtual void remove(const IdVector& ids) const = 0;
 
     virtual void write_index(
-        const tiledb::Context& ctx, const std::string& group_uri) const = 0;
+        const tiledb::Context& ctx,
+        const std::string& group_uri,
+        const std::string& storage_version) const = 0;
 
     [[nodiscard]] virtual size_t dimension() const = 0;
 
@@ -687,9 +691,11 @@ class IndexIVFFlat {
       //      index_.update(vectors_uri, ids, options);
     }
 
-    void write_index(const tiledb::Context& ctx, const std::string& group_uri)
-        const override {
-      impl_index_.write_index(ctx, group_uri);
+    void write_index(
+        const tiledb::Context& ctx,
+        const std::string& group_uri,
+        const std::string& storage_version) const override {
+      impl_index_.write_index(ctx, group_uri, storage_version);
     }
 
     // WIP

src/include/api/vamana_index.h

@@ -233,12 +233,14 @@ class IndexVamana {
   }
 
   void write_index(
-      const tiledb::Context& ctx, const std::string& group_uri) const {
+      const tiledb::Context& ctx,
+      const std::string& group_uri,
+      const std::string& storage_version = "") const {
     if (!index_) {
       throw std::runtime_error(
           "Cannot write_index() because there is no index.");
     }
-    index_->write_index(ctx, group_uri);
+    index_->write_index(ctx, group_uri, storage_version);
   }
 
   constexpr auto dimension() const {
@@ -287,7 +289,9 @@ class IndexVamana {
         std::optional<size_t> opt_L) = 0;
 
     virtual void write_index(
-        const tiledb::Context& ctx, const std::string& group_uri) const = 0;
+        const tiledb::Context& ctx,
+        const std::string& group_uri,
+        const std::string& storage_version) const = 0;
 
     [[nodiscard]] virtual size_t dimension() const = 0;
   };
@@ -390,9 +394,11 @@ class IndexVamana {
       }
     }
 
-    void write_index(const tiledb::Context& ctx, const std::string& group_uri)
-        const override {
-      impl_index_.write_index(ctx, group_uri);
+    void write_index(
+        const tiledb::Context& ctx,
+        const std::string& group_uri,
+        const std::string& storage_version) const override {
+      impl_index_.write_index(ctx, group_uri, storage_version);
     }
 
     size_t dimension() const override {

src/include/index/index_group.h

@@ -196,7 +196,7 @@ class base_index_group {
     metadata_.load_metadata(read_group);
     if (!empty(version_) && metadata_.storage_version_ != version_) {
       throw std::runtime_error(
-          "Version mismatch.  Requested " + version_ + " but found " +
+          "Version mismatch. Requested " + version_ + " but found " +
           metadata_.storage_version_);
     } else if (empty(version_)) {
       version_ = metadata_.storage_version_;

src/include/index/ivf_flat_index.h

@@ -737,14 +737,18 @@ class ivf_flat_index {
    * all of it to a TileDB group.  Since we have all of it in memory,
    * we write from the PartitionedMatrix base class.
    *
-   * @param group_uri
-   * @return bool indicating success or failure
+   * @param group_uri The URI of the TileDB group where the index will be saved
+   * @param storage_version The storage version to use. If empty, use the most
+   * defult version.
+   * @return Whether the write was successful
    */
   auto write_index(
-      const tiledb::Context& ctx, const std::string& group_uri) const {
+      const tiledb::Context& ctx,
+      const std::string& group_uri,
+      const std::string& storage_version = "") const {
     // Write the group
-    auto write_group =
-        ivf_flat_index_group(*this, ctx, group_uri, TILEDB_WRITE, timestamp_);
+    auto write_group = ivf_flat_index_group(
+        *this, ctx, group_uri, TILEDB_WRITE, timestamp_, storage_version);
 
     write_group.set_dimension(dimension_);
 

src/include/index/vamana_index.h

@@ -824,6 +824,8 @@ class vamana_index {
   /**
    * @brief Write the index to a TileDB group
    * @param group_uri The URI of the TileDB group where the index will be saved
+   * @param storage_version The storage version to use. If empty, use the most
+   * defult version.
    * @return Whether the write was successful
    *
    * The group consists of the original feature vectors, and the graph index,
@@ -836,12 +838,14 @@ class vamana_index {
    * the group?
    */
   auto write_index(
-      const tiledb::Context& ctx, const std::string& group_uri) const {
+      const tiledb::Context& ctx,
+      const std::string& group_uri,
+      const std::string& storage_version = "") const {
     // metadata: dimension, ntotal, L, R, B, alpha_min, alpha_max, medoid
     // Save as a group: metadata, feature_vectors, graph edges, offsets
 
-    auto write_group =
-        vamana_index_group(*this, ctx, group_uri, TILEDB_WRITE, timestamp_);
+    auto write_group = vamana_index_group(
+        *this, ctx, group_uri, TILEDB_WRITE, timestamp_, storage_version);
 
     // @todo Make this table-driven
     write_group.set_dimension(dimension_);

src/include/test/unit_api_vamana_index.cc

@@ -441,3 +441,60 @@ TEST_CASE("api_vamana_index: read index and query", "[api_vamana_index]") {
   auto recall = ((double)intersections_a) / ((double)nt * k_nn);
   CHECK(recall == 1.0);
 }
+
+TEST_CASE("api_vamana_index: storage_version", "[api_vamana_index]") {
+  auto ctx = tiledb::Context{};
+  using feature_type_type = uint8_t;
+  using id_type_type = uint32_t;
+  auto feature_type = "uint8";
+  auto id_type = "uint32";
+  auto adjacency_row_index_type = "uint32";
+  size_t dimensions = 3;
+
+  std::string index_uri =
+      (std::filesystem::temp_directory_path() / "api_vamana_index").string();
+  tiledb::VFS vfs(ctx);
+  if (vfs.is_dir(index_uri)) {
+    vfs.remove_dir(index_uri);
+  }
+
+  {
+    // First we create the index with a storage_version.
+    auto index = IndexVamana(std::make_optional<IndexOptions>(
+        {{"feature_type", feature_type},
+         {"id_type", id_type},
+         {"adjacency_row_index_type", adjacency_row_index_type}}));
+
+    size_t num_vectors = 0;
+    auto empty_training_vector_array =
+        FeatureVectorArray(dimensions, num_vectors, feature_type, id_type);
+    index.train(empty_training_vector_array);
+    index.add(empty_training_vector_array);
+    index.write_index(ctx, index_uri, "0.3");
+
+    CHECK(index.feature_type_string() == feature_type);
+    CHECK(index.id_type_string() == id_type);
+    CHECK(index.adjacency_row_index_type_string() == adjacency_row_index_type);
+  }
+
+  {
+    // Now make sure if we try to write it again with a different
+    // storage_version, we throw.
+    auto index = IndexVamana(ctx, index_uri);
+    auto training = ColMajorMatrixWithIds<feature_type_type, id_type_type>{
+        {{8, 6, 7}, {5, 3, 0}, {9, 5, 0}, {2, 7, 3}}, {10, 11, 12, 13}};
+
+    auto training_vector_array = FeatureVectorArray(training);
+    index.train(training_vector_array);
+    index.add(training_vector_array);
+
+    // Throw with the wrong version.
+    CHECK_THROWS_WITH(
+        index.write_index(ctx, index_uri, "0.4"),
+        "Version mismatch. Requested 0.4 but found 0.3");
+    // Succeed without a version.
+    index.write_index(ctx, index_uri);
+    // Succeed with the same version.
+    index.write_index(ctx, index_uri, "0.3");
+  }
+}

src/include/test/unit_flatpq_index.cc

@@ -829,6 +829,10 @@ TEST_CASE("flatpq_index: flatpq_index write and read", "[flatpq_index]") {
   tiledb::Context ctx;
   std::string flatpq_index_uri =
       (std::filesystem::temp_directory_path() / "tmp_flatpq_index").string();
+  tiledb::VFS vfs(ctx);
+  if (vfs.is_dir(flatpq_index_uri)) {
+    vfs.remove_dir(flatpq_index_uri);
+  }
   auto training_set =
       tdbColMajorMatrix<siftsmall_feature_type>(ctx, siftsmall_inputs_uri, 0);
   load(training_set);

src/include/test/unit_ivf_flat_group.cc

@@ -92,9 +92,6 @@ struct dummy_index {
   }
 };
 
-// The catch2 check for exception doesn't seem to be working correctly
-// @todo Fix this
-#if 0
 TEST_CASE(
     "ivf_flat_group: read constructor for non-existent group",
     "[ivf_flat_group]") {
@@ -104,7 +101,6 @@ TEST_CASE(
       ivf_flat_index_group(dummy_index{}, ctx, "I dont exist"),
       "Group uri I dont exist does not exist.");
 }
-#endif
 
 TEST_CASE("ivf_flat_group: write constructor - create", "[ivf_flat_group]") {
   std::string tmp_uri = (std::filesystem::temp_directory_path() /
@@ -119,7 +115,6 @@ TEST_CASE("ivf_flat_group: write constructor - create", "[ivf_flat_group]") {
 
   ivf_flat_index_group x =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
-  x.dump("Write constructor - create");
 }
 
 TEST_CASE(
@@ -137,11 +132,9 @@ TEST_CASE(
 
   ivf_flat_index_group x =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
-  x.dump("Write constructor - create before open");
 
   ivf_flat_index_group y =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
-  x.dump("Write constructor - open");
 }
 
 TEST_CASE(
@@ -159,11 +152,9 @@ TEST_CASE(
 
   ivf_flat_index_group x =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
-  x.dump("Write constructor - create before open");
 
   ivf_flat_index_group y =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_READ);
-  x.dump("Write constructor - open for read");
 }
 
 TEST_CASE(
@@ -182,15 +173,12 @@ TEST_CASE(
 
   ivf_flat_index_group x =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
-  x.dump("Write constructor - create before open");
 
   ivf_flat_index_group y =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
-  x.dump("Write constructor - open for write");
 
   ivf_flat_index_group z =
       ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_READ);
-  x.dump("Write constructor - open for read");
 }
 
 TEST_CASE(
@@ -365,3 +353,81 @@ TEST_CASE(
   CHECK(x.get_temp_size() == expected_temp_size + offset);
   CHECK(x.get_dimension() == expected_dimension + offset);
 }
+
+TEST_CASE("ivf_flat_group: storage version", "[ivf_flat_group]") {
+  std::string tmp_uri =
+      (std::filesystem::temp_directory_path() / "ivf_flat_group").string();
+
+  tiledb::Context ctx;
+  tiledb::VFS vfs(ctx);
+  if (vfs.is_dir(tmp_uri)) {
+    vfs.remove_dir(tmp_uri);
+  }
+
+  size_t expected_ingestion = 23094;
+  size_t expected_base = 9234;
+  size_t expected_partitions = 200;
+  size_t expected_temp_size = 11;
+  size_t expected_dimension = 19238;
+  auto offset = 2345;
+
+  ivf_flat_index_group x =
+      ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE);
+
+  SECTION("0.3") {
+    x = ivf_flat_index_group(
+        dummy_index{}, ctx, tmp_uri, TILEDB_WRITE, 0, "0.3");
+  }
+
+  SECTION("current_storage_version") {
+    x = ivf_flat_index_group(
+        dummy_index{}, ctx, tmp_uri, TILEDB_WRITE, 0, current_storage_version);
+  }
+
+  x.set_ingestion_timestamp(expected_ingestion + offset);
+  x.set_base_size(expected_base + offset);
+  x.set_num_partitions(expected_partitions + offset);
+  x.set_temp_size(expected_temp_size + offset);
+  x.set_dimension(expected_dimension + offset);
+
+  CHECK(size(x.get_all_ingestion_timestamps()) == 1);
+  CHECK(size(x.get_all_base_sizes()) == 1);
+  CHECK(size(x.get_all_num_partitions()) == 1);
+  CHECK(x.get_previous_ingestion_timestamp() == expected_ingestion + offset);
+  CHECK(x.get_previous_base_size() == expected_base + offset);
+  CHECK(x.get_previous_num_partitions() == expected_partitions + offset);
+  CHECK(x.get_temp_size() == expected_temp_size + offset);
+  CHECK(x.get_dimension() == expected_dimension + offset);
+}
+
+TEST_CASE("ivf_flat_group: invalid storage version", "[ivf_flat_group]") {
+  std::string tmp_uri =
+      (std::filesystem::temp_directory_path() / "ivf_flat_group").string();
+
+  tiledb::Context ctx;
+  tiledb::VFS vfs(ctx);
+  if (vfs.is_dir(tmp_uri)) {
+    vfs.remove_dir(tmp_uri);
+  }
+  CHECK_THROWS(ivf_flat_index_group(
+      dummy_index{}, ctx, tmp_uri, TILEDB_WRITE, 0, "invalid"));
+}
+
+TEST_CASE("ivf_flat_group: mismatched storage version", "[ivf_flat_group]") {
+  std::string tmp_uri =
+      (std::filesystem::temp_directory_path() / "ivf_flat_group").string();
+
+  tiledb::Context ctx;
+  tiledb::VFS vfs(ctx);
+  if (vfs.is_dir(tmp_uri)) {
+    vfs.remove_dir(tmp_uri);
+  }
+
+  ivf_flat_index_group x =
+      ivf_flat_index_group(dummy_index{}, ctx, tmp_uri, TILEDB_WRITE, 0, "0.3");
+
+  CHECK_THROWS_WITH(
+      ivf_flat_index_group(
+          dummy_index{}, ctx, tmp_uri, TILEDB_WRITE, 0, "different_version"),
+      "Version mismatch. Requested different_version but found 0.3");
+}