Vamana type-erased index can be opened at a specific timestamp (#354)

Author: jparismorgan

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

@@ -34,8 +34,8 @@ def __init__(
         timestamp=None,
         **kwargs,
     ):
-        super().__init__(uri=uri, config=config, timestamp=timestamp)
         self.index_type = INDEX_TYPE
+        super().__init__(uri=uri, config=config, timestamp=timestamp)
         self._index = None
         self.db_uri = self.group[
             storage_formats[self.storage_version]["PARTS_ARRAY_NAME"]

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

@@ -8,6 +8,7 @@
 from tiledb.vector_search.module import *
 from tiledb.vector_search.storage_formats import storage_formats
 from tiledb.vector_search.utils import add_to_group
+from tiledb.vector_search.utils import is_type_erased_index
 
 MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 MAX_INT32 = np.iinfo(np.dtype("int32")).max
@@ -466,6 +467,7 @@ def clear_history(
     ):
         with tiledb.scope_ctx(ctx_or_config=config):
             group = tiledb.Group(uri, "r")
+            index_type = group.meta.get("index_type", "")
             storage_version = group.meta.get("storage_version", "0.1")
             if not storage_formats[storage_version]["SUPPORT_TIMETRAVEL"]:
                 raise ValueError(
@@ -490,7 +492,9 @@ def clear_history(
                 if ingestion_timestamp > timestamp:
                     new_ingestion_timestamps.append(ingestion_timestamp)
                     new_base_sizes.append(base_sizes[i])
-                    new_partition_history.append(partition_history[i])
+                    # Type erased indexes don't have partition_history, skip to avoid crash.
+                    if not is_type_erased_index(index_type):
+                        new_partition_history.append(partition_history[i])
                 i += 1
             if len(new_ingestion_timestamps) == 0:
                 new_ingestion_timestamps = [0]
@@ -502,7 +506,9 @@ def clear_history(
             group = tiledb.Group(uri, "w")
             group.meta["ingestion_timestamps"] = json.dumps(new_ingestion_timestamps)
             group.meta["base_sizes"] = json.dumps(new_base_sizes)
-            group.meta["partition_history"] = json.dumps(new_partition_history)
+            # Type erased indexes don't have partition_history, skip to avoid crash.
+            if not is_type_erased_index(index_type):
+                group.meta["partition_history"] = json.dumps(new_partition_history)
             group.close()
 
             group = tiledb.Group(uri, "r")

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

@@ -1632,7 +1632,13 @@ def ingest_vamana(
 
         ctx = vspy.Ctx(config)
         index = vspy.IndexVamana(ctx, index_group_uri)
-        data = vspy.FeatureVectorArray(ctx, parts_array_uri, ids_array_uri)
+        data = vspy.FeatureVectorArray(
+            ctx,
+            parts_array_uri,
+            ids_array_uri,
+            0,
+            index_timestamp if index_timestamp else 0,
+        )
         index.train(data)
         index.add(data)
         index.write_index(ctx, index_group_uri, index_timestamp)

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

@@ -48,8 +48,8 @@ def __init__(
         memory_budget: int = -1,
         **kwargs,
     ):
-        super().__init__(uri=uri, config=config, timestamp=timestamp)
         self.index_type = INDEX_TYPE
+        super().__init__(uri=uri, config=config, timestamp=timestamp)
         self.db_uri = self.group[
             storage_formats[self.storage_version]["PARTS_ARRAY_NAME"]
             + self.index_version

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

@@ -173,12 +173,22 @@ void init_type_erased_module(py::module_& m) {
           py::keep_alive<1, 2>()  // FeatureVectorArray should keep ctx alive.
           )
       .def(
-          py::init<
-              const tiledb::Context&,
-              const std::string&,
-              const std::string&>(),
-          py::keep_alive<1, 2>()  // FeatureVectorArray should keep ctx alive.
-          )
+          "__init__",
+          [](FeatureVectorArray& instance,
+             const tiledb::Context& ctx,
+             const std::string& uri,
+             const std::string& ids_uri,
+             size_t num_vectors,
+             size_t timestamp) {
+            new (&instance)
+                FeatureVectorArray(ctx, uri, ids_uri, num_vectors, timestamp);
+          },
+          py::keep_alive<1, 2>(),  // FeatureVectorArray should keep ctx alive.
+          py::arg("ctx"),
+          py::arg("uri"),
+          py::arg("ids_uri") = "",
+          py::arg("num_vectors") = 0,
+          py::arg("timestamp") = 0)
       .def(py::init<size_t, size_t, const std::string&, const std::string&>())
       .def("dimension", &FeatureVectorArray::dimension)
       .def("num_vectors", &FeatureVectorArray::num_vectors)
@@ -261,9 +271,17 @@ void init_type_erased_module(py::module_& m) {
 
   py::class_<IndexVamana>(m, "IndexVamana")
       .def(
-          py::init<const tiledb::Context&, const std::string&>(),
-          py::keep_alive<1, 2>()  // IndexVamana should keep ctx alive.
-          )
+          "__init__",
+          [](IndexVamana& instance,
+             const tiledb::Context& ctx,
+             const std::string& group_uri,
+             size_t timestamp) {
+            new (&instance) IndexVamana(ctx, group_uri, timestamp);
+          },
+          py::keep_alive<1, 2>(),  // IndexVamana should keep ctx alive.
+          py::arg("ctx"),
+          py::arg("group_uri"),
+          py::arg("timestamp") = 0)
       .def(
           "__init__",
           [](IndexVamana& instance, py::kwargs kwargs) {

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

@@ -35,7 +35,12 @@ def __init__(
     ):
         super().__init__(uri=uri, config=config, timestamp=timestamp)
         self.index_type = INDEX_TYPE
-        self.index = vspy.IndexVamana(vspy.Ctx(config), uri)
+        # TODO(paris): Support (start, end) timestamps and remove.
+        type_erased_timestamp = timestamp
+        if isinstance(timestamp, tuple):
+            type_erased_timestamp = timestamp[1]
+        type_erased_timestamp = type_erased_timestamp if type_erased_timestamp else 0
+        self.index = vspy.IndexVamana(self.ctx, uri, type_erased_timestamp)
         self.db_uri = self.group[
             storage_formats[self.storage_version]["PARTS_ARRAY_NAME"]
         ].uri

apis/python/test/test_ingestion.py

@@ -560,11 +560,6 @@ def test_ingestion_with_batch_updates(tmp_path):
     gt_i, gt_d = get_groundtruth(dataset_dir, k)
 
     for index_type, index_class in zip(INDEXES, INDEX_CLASSES):
-        # TODO(paris): Fix Vamana bug and re-enable:
-        # tiledb.cc.TileDBError: [TileDB::ArrayDirectory] Error: Cannot open array; Array does not exist.
-        if index_type == "VAMANA":
-            continue
-
         index_uri = os.path.join(tmp_path, f"array_{index_type}")
         index = ingest(
             index_type=index_type,
@@ -665,10 +660,6 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
         index_uri = move_local_index_to_new_location(index_uri)
         index = index_class(uri=index_uri, timestamp=(0, 101))
         _, result = index.query(queries, k=k, nprobe=partitions)
-        # TODO(paris): Fix Vamana accuracy bug and re-enable:
-        # assert 0.105 == 1.0
-        if index_type == "VAMANA":
-            continue
         assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
         index = index_class(uri=index_uri, timestamp=(0, None))
         _, result = index.query(queries, k=k, nprobe=partitions)
@@ -773,6 +764,9 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
         Index.clear_history(
             uri=index_uri, timestamp=index.latest_ingestion_timestamp - 1
         )
+        if index_type == "VAMANA":
+            # TODO(paris): Re-enable once we support (start, end) timestamps for Vamana.
+            continue
         index = index_class(uri=index_uri, timestamp=1)
         _, result = index.query(queries, k=k, nprobe=partitions)
         assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0

apis/python/test/test_type_erased_module.py

@@ -338,4 +338,4 @@ def test_inplace_build_infinite_query_IndexIVFFlat():
         if nprobe == 8:
             assert recall > 0.925
         if nprobe == 32:
-            assert recall >= 0.999
+            assert recall >= 0.998