Add support for ingest() with index_timestamp for Vamana type-erased index (#341)

Author: jparismorgan

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

@@ -90,7 +90,7 @@ def ingest(
     updates_uri: str
         Updates
     index_timestamp: int
-        Timestamp to use for writing and reading data. By default it sues the current unix ms timestamp.
+        Timestamp to use for writing and reading data. By default it uses the current unix ms timestamp.
     config: None
         config dictionary, defaults to None
     namespace: str
@@ -1633,7 +1633,7 @@ def ingest_vamana(
             data = vspy.FeatureVectorArray(ctx, parts_array_uri, ids_array_uri)
             index.train(data)
             index.add(data)
-            index.write_index(ctx, index_group_uri)
+            index.write_index(ctx, index_group_uri, index_timestamp)
 
     def write_centroids(
         centroids: np.ndarray,

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

@@ -141,6 +141,7 @@ def query_internal(
         resources: Optional[Mapping[str, Any]] = None,
         num_partitions: int = -1,
         num_workers: int = -1,
+        **kwargs,
     ):
         """
         Query an IVF_FLAT index

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

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

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

@@ -121,5 +121,5 @@ def create(
     )
     index.train(empty_vector)
     index.add(empty_vector)
-    index.write_index(ctx, uri, storage_version)
+    index.write_index(ctx, uri, 0, storage_version)
     return VamanaIndex(uri=uri, config=config, memory_budget=1000000)

apis/python/test/test_index.py

@@ -1,5 +1,4 @@
 import json
-import time
 
 import numpy as np
 import pytest
@@ -57,31 +56,20 @@ def check_default_metadata(
     assert type(group.meta["index_type"]) == str
 
     assert "base_sizes" in group.meta
-    if is_type_erased_index(expected_index_type):
-        # NOTE(paris): Type-erased indexes have two values upon creation.
-        assert group.meta["base_sizes"] == "[0,0]"
-    else:
-        assert group.meta["base_sizes"] == json.dumps([0])
+    assert group.meta["base_sizes"] == json.dumps([0])
     assert type(group.meta["base_sizes"]) == str
 
     assert "ingestion_timestamps" in group.meta
-    if is_type_erased_index(expected_index_type):
-        # NOTE(paris): Type-erased indexes have two values upon creation.
-        ingestion_timestamps = json.loads(group.meta["ingestion_timestamps"])
-        assert len(ingestion_timestamps) == 2
-        assert ingestion_timestamps[0] == 0
-        current_time_ms = int(time.time() * 1000)
-        assert ingestion_timestamps[1] < current_time_ms
-        assert ingestion_timestamps[1] > current_time_ms - 1000 * 5
-    else:
-        assert group.meta["ingestion_timestamps"] == json.dumps([0])
+    assert group.meta["ingestion_timestamps"] == json.dumps([0])
     assert type(group.meta["ingestion_timestamps"]) == str
 
     if not is_type_erased_index(expected_index_type):
         # NOTE(paris): Type-erased indexes do not write has_updates.
         assert "has_updates" in group.meta
         assert group.meta["has_updates"] == 0
         assert type(group.meta["has_updates"]) == np.int64
+    else:
+        assert "has_updates" not in group.meta
 
 
 def test_flat_index(tmp_path):
@@ -279,7 +267,7 @@ def test_vamana_index(tmp_path):
 
 def test_delete_invalid_index(tmp_path):
     # We don't throw with an invalid uri.
-    Index.delete_index(uri="invalid_uri", config=tiledb.cloud.Config())
+    Index.delete_index(uri="invalid_uri", config={})
 
 
 def test_delete_index(tmp_path):
@@ -289,7 +277,7 @@ def test_delete_index(tmp_path):
     for index_type, index_class in zip(indexes, index_classes):
         index_uri = os.path.join(tmp_path, f"array_{index_type}")
         ingest(index_type=index_type, index_uri=index_uri, input_vectors=data)
-        Index.delete_index(uri=index_uri, config=tiledb.cloud.Config())
+        Index.delete_index(uri=index_uri, config={})
         with pytest.raises(tiledb.TileDBError) as error:
             index_class(uri=index_uri)
         assert "does not exist" in str(error.value)

apis/python/test/test_ingestion.py

@@ -607,11 +607,6 @@ def test_ingestion_with_updates_and_timetravel(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:
-        # ValueError: New ingestion timestamp: 1 can't be smaller that the latest ingestion timestamp: 1713444057062
-        if index_type == "VAMANA":
-            continue
-
         index_uri = os.path.join(tmp_path, f"array_{index_type}")
         index = ingest(
             index_type=index_type,
@@ -647,6 +642,11 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
         index = index_class(uri=index_uri, timestamp=101)
         _, result = index.query(queries, k=k, nprobe=partitions)
         assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+
+        # TODO(paris): Fix Vamana bug and re-enable:
+        if index_type == "VAMANA":
+            continue
+
         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)
@@ -837,11 +837,6 @@ def test_ingestion_with_additions_and_timetravel(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:
-        # ValueError: New ingestion timestamp: 1 can't be smaller that the latest ingestion timestamp: 1713444057062
-        if index_type == "VAMANA":
-            continue
-
         index_uri = os.path.join(tmp_path, f"array_{index_type}")
         index = ingest(
             index_type=index_type,
@@ -865,14 +860,17 @@ def test_ingestion_with_additions_and_timetravel(tmp_path):
             )
             updated_ids[i] = i + update_ids_offset
 
-        index_uri = move_local_index_to_new_location(index_uri)
+        # 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":
+            index_uri = move_local_index_to_new_location(index_uri)
         index = index_class(uri=index_uri)
-        _, result = index.query(queries, k=k, nprobe=partitions)
-        assert 0.45 < accuracy(result, gt_i) < 0.55
+        _, result = index.query(queries, k=k, nprobe=partitions, opt_l=k * 2)
+        assert 0.45 < accuracy(result, gt_i)
 
         index = index.consolidate_updates()
-        _, result = index.query(queries, k=k, nprobe=partitions)
-        assert 0.45 < accuracy(result, gt_i) < 0.55
+        _, result = index.query(queries, k=k, nprobe=partitions, opt_l=k * 2)
+        assert 0.45 < accuracy(result, gt_i)
 
 
 def test_ivf_flat_ingestion_tdb_random_sampling_policy(tmp_path):

src/include/api/vamana_index.h

@@ -235,12 +235,13 @@ class IndexVamana {
   void write_index(
       const tiledb::Context& ctx,
       const std::string& group_uri,
-      const std::string& storage_version = "") const {
+      std::optional<size_t> timestamp = std::nullopt,
+      const std::string& storage_version = "") {
     if (!index_) {
       throw std::runtime_error(
           "Cannot write_index() because there is no index.");
     }
-    index_->write_index(ctx, group_uri, storage_version);
+    index_->write_index(ctx, group_uri, timestamp, storage_version);
   }
 
   constexpr auto dimension() const {
@@ -291,7 +292,8 @@ class IndexVamana {
     virtual void write_index(
         const tiledb::Context& ctx,
         const std::string& group_uri,
-        const std::string& storage_version) const = 0;
+        std::optional<size_t> timestamp,
+        const std::string& storage_version) = 0;
 
     [[nodiscard]] virtual size_t dimension() const = 0;
   };
@@ -397,8 +399,9 @@ class IndexVamana {
     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);
+        std::optional<size_t> timestamp,
+        const std::string& storage_version) override {
+      impl_index_.write_index(ctx, group_uri, timestamp, storage_version);
     }
 
     size_t dimension() const override {

src/include/index/index_group.h

@@ -239,7 +239,8 @@ class base_index_group {
     if (exists(cached_ctx_)) {
       /** Load the current group metadata */
       init_for_open(cfg);
-      if (index_timestamp_ < metadata_.ingestion_timestamps_.back()) {
+      if (!metadata_.ingestion_timestamps_.empty() &&
+          index_timestamp_ < metadata_.ingestion_timestamps_.back()) {
         throw std::runtime_error(
             "Requested write timestamp " + std::to_string(index_timestamp_) +
             " is not greater than " +
@@ -373,9 +374,6 @@ class base_index_group {
   auto get_previous_ingestion_timestamp() const {
     return metadata_.ingestion_timestamps_.back();
   }
-  auto get_ingestion_timestamp() const {
-    return metadata_.ingestion_timestamps_[timetravel_index_];
-  }
   auto append_ingestion_timestamp(size_t timestamp) {
     metadata_.ingestion_timestamps_.push_back(timestamp);
   }

src/include/index/vamana_group.h

@@ -257,9 +257,9 @@ class vamana_index_group : public base_index_group<vamana_index_group<Index>> {
     metadata_.adjacency_row_index_type_str_ =
         type_to_string_v<typename index_type::adjacency_row_index_type>;
 
-    metadata_.ingestion_timestamps_ = {0};
-    metadata_.base_sizes_ = {0};
-    metadata_.num_edges_history_ = {0};
+    metadata_.ingestion_timestamps_ = {};
+    metadata_.base_sizes_ = {};
+    metadata_.num_edges_history_ = {};
     metadata_.temp_size_ = 0;
     metadata_.dimension_ = this->get_dimension();
 

src/include/index/vamana_index.h

@@ -840,7 +840,11 @@ class vamana_index {
   auto write_index(
       const tiledb::Context& ctx,
       const std::string& group_uri,
-      const std::string& storage_version = "") const {
+      std::optional<size_t> timestamp = std::nullopt,
+      const std::string& storage_version = "") {
+    if (timestamp.has_value()) {
+      timestamp_ = timestamp.value();
+    }
     // metadata: dimension, ntotal, L, R, B, alpha_min, alpha_max, medoid
     // Save as a group: metadata, feature_vectors, graph edges, offsets