Fix TileDB URIs with type-erased indexes, add Vamana test to test_cloud.py (#330)

Author: jparismorgan

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

@@ -8,6 +8,7 @@
 from tiledb.vector_search.module import *
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
 from tiledb.vector_search.storage_formats import storage_formats
+from tiledb.vector_search.storage_formats import validate_storage_version
 
 MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 INDEX_TYPE = "VAMANA"
@@ -106,6 +107,7 @@ def create(
     **kwargs,
 ) -> VamanaIndex:
     warnings.warn("The Vamana index is not yet supported, please use with caution.")
+    validate_storage_version(storage_version)
     ctx = vspy.Ctx(config)
     index = vspy.IndexVamana(
         feature_type=np.dtype(vector_type).name,

apis/python/test/test_cloud.py

@@ -27,6 +27,7 @@ def setUpClass(cls):
         rand_name = random_name("vector_search")
         test_path = f"tiledb://{namespace}/{storage_path}/{rand_name}"
         cls.flat_index_uri = f"{test_path}/test_flat_array"
+        cls.vamana_index_uri = f"{test_path}/vamana_array"
         cls.ivf_flat_index_uri = f"{test_path}/test_ivf_flat_array"
         cls.ivf_flat_random_sampling_index_uri = (
             f"{test_path}/test_ivf_flat_random_sampling_array"
@@ -35,6 +36,7 @@ def setUpClass(cls):
     @classmethod
     def tearDownClass(cls):
         vs.Index.delete_index(uri=cls.flat_index_uri, config=tiledb.cloud.Config())
+        vs.Index.delete_index(uri=cls.vamana_index_uri, config=tiledb.cloud.Config())
         vs.Index.delete_index(uri=cls.ivf_flat_index_uri, config=tiledb.cloud.Config())
         vs.Index.delete_index(
             uri=cls.ivf_flat_random_sampling_index_uri, config=tiledb.cloud.Config()
@@ -70,6 +72,31 @@ def test_cloud_flat(self):
         _, result_i = index.query(queries, k=k)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
+    def test_cloud_vamana(self):
+        source_uri = "tiledb://TileDB-Inc/sift_10k"
+        queries_uri = siftsmall_query_file
+        gt_uri = siftsmall_groundtruth_file
+        index_uri = CloudTests.vamana_index_uri
+        k = 100
+        nqueries = 100
+
+        load_fvecs(queries_uri)
+        gt_i, gt_d = get_groundtruth_ivec(gt_uri, k=k, nqueries=nqueries)
+
+        vs.ingest(
+            index_type="VAMANA",
+            index_uri=index_uri,
+            source_uri=source_uri,
+            input_vectors_per_work_item=5000,
+            config=tiledb.cloud.Config().dict(),
+            # TODO(paris): Fix and then change to Mode.BATCH.
+            mode=Mode.LOCAL,
+        )
+
+        # TODO(paris): Fix error from loading this URI and then re-enable, and add the rest of the test.
+        # tiledb_index_uri = groups.info(index_uri).tiledb_uri
+        # vs.vamana_index.VamanaIndex(uri=tiledb_index_uri)
+
     def test_cloud_ivf_flat(self):
         source_uri = "tiledb://TileDB-Inc/sift_10k"
         queries_uri = siftsmall_query_file

src/include/detail/linalg/tdb_helpers.h

@@ -37,6 +37,7 @@
 #include <tiledb/array.h>
 #include <tiledb/tiledb>
 #include "stats.h"
+#include "tiledb/group_experimental.h"
 
 namespace tiledb_helpers {
 
@@ -81,6 +82,19 @@ inline void submit_query(
   query.submit();
 }
 
+// Adds an object to a group. Automatically infers whether to use a relative
+// path or absolute path. NOTE(paris): We use absolute paths for tileDB URIs
+// because of a bug tracked in SC39197, once that is fixed everything can use
+// relative paths.
+inline void add_to_group(
+    tiledb::Group& group, const std::string& uri, const std::string& name) {
+  if (uri.find("tiledb://") == 0) {
+    group.add_member(uri, false, name);
+  } else {
+    group.add_member(name, true, name);
+  }
+}
+
 }  // namespace tiledb_helpers
 
 #endif

src/include/index/flatpq_index.h

@@ -653,11 +653,11 @@ class flatpq_index {
 
     auto centroids_uri = group_uri + "/centroids";
     write_matrix(ctx, centroids_, centroids_uri);
-    write_group.add_member("centroids", true, "centroids");
+    tiledb_helpers::add_to_group(write_group, centroids_uri, "centroids");
 
     auto pq_vectors_uri = group_uri + "/pq_vectors";
     write_matrix(ctx, pq_vectors_, pq_vectors_uri);
-    write_group.add_member("pq_vectors", true, "pq_vectors");
+    tiledb_helpers::add_to_group(write_group, pq_vectors_uri, "pq_vectors");
 
     for (size_t subspace = 0; subspace < num_subspaces_; ++subspace) {
       std::ostringstream oss;
@@ -666,8 +666,8 @@ class flatpq_index {
 
       auto distance_table_uri = group_uri + "/distance_table_" + number;
       write_matrix(ctx, distance_tables_[subspace], distance_table_uri);
-      write_group.add_member(
-          "distance_table_" + number, true, "distance_table_" + number);
+      tiledb_helpers::add_to_group(
+          write_group, distance_table_uri, "distance_table_" + number);
     }
     write_group.close();
     return true;

src/include/index/index_group.h

@@ -182,8 +182,7 @@ class base_index_group {
    * @param ctx
    */
   void init_for_open(const tiledb::Config& cfg) {
-    tiledb::VFS vfs(cached_ctx_);
-    if (!vfs.is_dir(group_uri_)) {
+    if (!exists(cached_ctx_)) {
       throw std::runtime_error(
           "Group uri " + std::string(group_uri_) + " does not exist.");
     }
@@ -264,9 +263,7 @@ class base_index_group {
    * @param version
    */
   void open_for_write(const tiledb::Config& cfg) {
-    tiledb::VFS vfs(cached_ctx_);
-
-    if (vfs.is_dir(group_uri_)) {
+    if (exists(cached_ctx_)) {
       /** Load the current group metadata */
       init_for_open(cfg);
       if (index_timestamp_ < metadata_.ingestion_timestamps_.back()) {

src/include/index/ivf_flat_group.h

@@ -197,9 +197,8 @@ class ivf_flat_index_group
         this->get_dimension(),
         default_tile_extent,
         default_compression);
-    // write_group.add_member(centroids_uri(), true, centroids_array_name());
-    write_group.add_member(
-        centroids_array_name(), true, centroids_array_name());
+    tiledb_helpers::add_to_group(
+        write_group, centroids_uri(), centroids_array_name());
 
     create_empty_for_matrix<
         typename index_type::feature_type,
@@ -211,22 +210,20 @@ class ivf_flat_index_group
         this->get_dimension(),
         default_tile_extent,
         default_compression);
-    // write_group.add_member(parts_uri(), true, parts_array_name());
-    write_group.add_member(parts_array_name(), true, parts_array_name());
+    tiledb_helpers::add_to_group(write_group, parts_uri(), parts_array_name());
 
     create_empty_for_vector<typename index_type::id_type>(
         cached_ctx_, ids_uri(), default_domain, tile_size, default_compression);
-    // write_group.add_member(ids_uri(), true, ids_array_name());
-    write_group.add_member(ids_array_name(), true, ids_array_name());
+    tiledb_helpers::add_to_group(write_group, ids_uri(), ids_array_name());
 
     create_empty_for_vector<typename index_type::indices_type>(
         cached_ctx_,
         indices_uri(),
         default_domain,
         default_tile_extent,
         default_compression);
-    // write_group.add_member(indices_uri(), true, indices_array_name());
-    write_group.add_member(indices_array_name(), true, indices_array_name());
+    tiledb_helpers::add_to_group(
+        write_group, indices_uri(), indices_array_name());
 
     // Store the metadata if all of the arrays were created successfully
     metadata_.store_metadata(write_group);

src/include/index/ivf_flat_index.h

@@ -792,8 +792,6 @@ class ivf_flat_index {
       const std::string& parts_uri,
       const std::string& ids_uri,
       const std::string& indices_uri) const {
-    tiledb::VFS vfs(ctx);
-
     write_matrix(ctx, centroids_, centroids_uri, 0, true);
     write_matrix(ctx, *partitioned_vectors_, parts_uri, 0, true);
     write_vector(ctx, partitioned_vectors_->ids(), ids_uri, 0, true);

src/include/index/vamana_group.h

@@ -32,6 +32,7 @@
 #ifndef TILEDB_VAMANA_GROUP_H
 #define TILEDB_VAMANA_GROUP_H
 
+#include "detail/linalg/tdb_helpers.h"
 #include "index/index_defs.h"
 #include "index/index_group.h"
 #include "index/vamana_metadata.h"
@@ -280,45 +281,45 @@ class vamana_index_group : public base_index_group<vamana_index_group<Index>> {
         this->get_dimension(),
         default_tile_extent,
         default_compression);
-    write_group.add_member(
-        feature_vectors_array_name(), true, feature_vectors_array_name());
+    tiledb_helpers::add_to_group(
+        write_group, feature_vectors_uri(), feature_vectors_array_name());
 
     create_empty_for_vector<typename index_type::id_type>(
         cached_ctx_,
         feature_vector_ids_uri(),
         default_domain,
         tile_size,
         default_compression);
-    write_group.add_member(
-        feature_vector_ids_name(), true, feature_vector_ids_name());
+    tiledb_helpers::add_to_group(
+        write_group, feature_vector_ids_uri(), feature_vector_ids_name());
 
     create_empty_for_vector<typename index_type::score_type>(
         cached_ctx_,
         adjacency_scores_uri(),
         default_domain,
         tile_size,
         default_compression);
-    write_group.add_member(
-        adjacency_scores_array_name(), true, adjacency_scores_array_name());
+    tiledb_helpers::add_to_group(
+        write_group, adjacency_scores_uri(), adjacency_scores_array_name());
 
     create_empty_for_vector<typename index_type::id_type>(
         cached_ctx_,
         adjacency_ids_uri(),
         default_domain,
         tile_size,
         default_compression);
-    write_group.add_member(
-        adjacency_ids_array_name(), true, adjacency_ids_array_name());
+    tiledb_helpers::add_to_group(
+        write_group, adjacency_ids_uri(), adjacency_ids_array_name());
 
     create_empty_for_vector<typename index_type::id_type>(
         cached_ctx_,
         adjacency_row_index_uri(),
         default_domain,
         tile_size,
         default_compression);
-    write_group.add_member(
-        adjacency_row_index_array_name(),
-        true,
+    tiledb_helpers::add_to_group(
+        write_group,
+        adjacency_row_index_uri(),
         adjacency_row_index_array_name());
 
     // Store the metadata if all of the arrays were created successfully

src/include/tdb_defs.h

@@ -90,7 +90,7 @@ constexpr auto type_to_tiledb_v = tiledb::impl::type_to_tiledb<T>::tiledb_type;
   if (str == "uint64") {
     return TILEDB_UINT64;
   }
-  throw std::runtime_error("Unsupported datatype");
+  throw std::runtime_error("Unsupported datatype: " + str);
 }
 
 // cf type_to_str(tiledb_datatype_t type) in tiledb/type.h