Add more C++ and Python temporal policy unit tests (#359)

Author: jparismorgan

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

@@ -314,7 +314,7 @@ void init_type_erased_module(py::module_& m) {
             new (&instance) IndexVamana(
                 ctx,
                 group_uri,
-                timestamp == 0 ? TemporalPolicy() :
+                timestamp == 0 ? TemporalPolicy(TimeTravel, 0) :
                                  TemporalPolicy(TimeTravel, timestamp));
           },
           py::keep_alive<1, 2>(),  // IndexVamana should keep ctx alive.

apis/python/test/common.py

@@ -1,3 +1,4 @@
+import json
 import os
 import random
 import shutil
@@ -392,3 +393,13 @@ def delete_uri(uri, config):
             else:
                 raise err
         group.delete(recursive=True)
+
+
+def load_metadata(index_uri):
+    group = tiledb.Group(index_uri, "r")
+    ingestion_timestamps = [
+        int(x) for x in list(json.loads(group.meta.get("ingestion_timestamps", "[]")))
+    ]
+    base_sizes = [int(x) for x in list(json.loads(group.meta.get("base_sizes", "[]")))]
+    group.close()
+    return ingestion_timestamps, base_sizes

apis/python/test/test_index.py

@@ -1,9 +1,11 @@
 import json
+import time
 
 import numpy as np
 import pytest
 from array_paths import *
 from common import *
+from common import load_metadata
 
 import tiledb.vector_search.index as ind
 from tiledb.vector_search import Index
@@ -125,6 +127,11 @@ def test_ivf_flat_index(tmp_path):
     index = ivf_flat_index.create(
         uri=uri, dimensions=3, vector_type=vector_type, partitions=partitions
     )
+
+    ingestion_timestamps, base_sizes = load_metadata(uri)
+    assert base_sizes == [0]
+    assert ingestion_timestamps == [0]
+
     query_and_check(
         index,
         np.array([[2, 2, 2]], dtype=np.float32),
@@ -147,6 +154,13 @@ def test_ivf_flat_index(tmp_path):
     )
 
     index = index.consolidate_updates()
+    # TODO(SC-46771): Investigate whether we should overwrite the existing metadata during the first
+    # ingestion of Python indexes. I believe as it's currently written we have a bug here.
+    # ingestion_timestamps, base_sizes = load_metadata(uri)
+    # assert base_sizes == [5]
+    # timestamp_5_minutes_from_now = int((time.time() + 5 * 60) * 1000)
+    # timestamp_5_minutes_ago = int((time.time() - 5 * 60) * 1000)
+    # assert ingestion_timestamps[0] > timestamp_5_minutes_ago and ingestion_timestamps[0] < timestamp_5_minutes_from_now
 
     query_and_check(
         index, np.array([[2, 2, 2]], dtype=np.float32), 3, {1, 2, 3}, nprobe=partitions
@@ -224,6 +238,10 @@ def test_vamana_index(tmp_path):
         vector_type=np.dtype(vector_type),
     )
 
+    ingestion_timestamps, base_sizes = load_metadata(uri)
+    assert base_sizes == [0]
+    assert ingestion_timestamps == [0]
+
     queries = np.array([[2, 2, 2]], dtype=np.float32)
     distances, ids = index.query(queries, k=1)
     assert distances.shape == (1, 1)
@@ -251,6 +269,16 @@ def test_vamana_index(tmp_path):
 
     index = index.consolidate_updates()
 
+    # During the first ingestion we overwrite the metadata and end up with a single base size and ingestion timestamp.
+    ingestion_timestamps, base_sizes = load_metadata(uri)
+    assert base_sizes == [5]
+    timestamp_5_minutes_from_now = int((time.time() + 5 * 60) * 1000)
+    timestamp_5_minutes_ago = int((time.time() - 5 * 60) * 1000)
+    assert (
+        ingestion_timestamps[0] > timestamp_5_minutes_ago
+        and ingestion_timestamps[0] < timestamp_5_minutes_from_now
+    )
+
     # Check that we throw if we query with an invalid opt_l.
     with pytest.raises(ValueError):
         index.query(queries, k=3, opt_l=2)

apis/python/test/test_ingestion.py

@@ -4,6 +4,7 @@
 import pytest
 from array_paths import *
 from common import *
+from common import load_metadata
 
 from tiledb.cloud.dag import Mode
 from tiledb.vector_search.flat_index import FlatIndex
@@ -513,16 +514,16 @@ def test_ingestion_with_updates(tmp_path):
             partitions=partitions,
         )
 
-        # TODO(paris): Fix Vamana to have same metadata as Python and re-enable.
-        # with tiledb.Group(index_uri, "r", ctx={}) as group:
-        #     ingestion_timestamps = [int(x) for x in list(json.loads(group.meta.get("ingestion_timestamps", "[]")))]
-        #     base_sizes = [int(x) for x in list(json.loads(group.meta.get("base_sizes", "[]")))]
-        #     assert len(ingestion_timestamps) == 1
-        #     assert len(base_sizes) == 1
-        #     assert base_sizes[0] == 1000
-        #     timestamp_2030 = 1903946089000
-        #     timestamp_5_minutes_ago = int((time.time() - 5 * 60) * 1000)
-        #     assert ingestion_timestamps[0] > timestamp_5_minutes_ago and ingestion_timestamps[0] < timestamp_2030
+        ingestion_timestamps, base_sizes = load_metadata(index_uri)
+        assert base_sizes == [1000]
+        assert len(ingestion_timestamps) == 1
+        timestamp_5_minutes_from_now = int((time.time() + 5 * 60) * 1000)
+        timestamp_5_minutes_ago = int((time.time() - 5 * 60) * 1000)
+        assert (
+            ingestion_timestamps[0] > timestamp_5_minutes_ago
+            and ingestion_timestamps[0] < timestamp_5_minutes_from_now
+        )
+        ingestion_timestamp = ingestion_timestamps[0]
 
         _, result = index.query(queries, k=k, nprobe=nprobe)
         assert accuracy(result, gt_i) == 1.0
@@ -548,6 +549,16 @@ def test_ingestion_with_updates(tmp_path):
         _, result = index.query(queries, k=k, nprobe=20)
         assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
 
+        ingestion_timestamps, base_sizes = load_metadata(index_uri)
+        assert base_sizes == [1000, 1000]
+        assert len(ingestion_timestamps) == 2
+        assert ingestion_timestamps[0] == ingestion_timestamp
+        assert (
+            ingestion_timestamps[1] != ingestion_timestamp
+            and ingestion_timestamps[1] > timestamp_5_minutes_ago
+            and ingestion_timestamps[1] < timestamp_5_minutes_from_now
+        )
+
         assert vfs.dir_size(index_uri) > 0
         Index.delete_index(uri=index_uri, config={})
         assert vfs.dir_size(index_uri) == 0
@@ -643,6 +654,10 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
             index_timestamp=1,
         )
 
+        ingestion_timestamps, base_sizes = load_metadata(index_uri)
+        assert ingestion_timestamps == [1]
+        assert base_sizes == [1000]
+
         if index_type == "IVF_FLAT":
             assert index.partitions == partitions
 
@@ -651,7 +666,8 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
 
         update_ids_offset = MAX_UINT64 - size
         updated_ids = {}
-        for i in range(2, 102):
+        timestamp_end = 102
+        for i in range(2, timestamp_end):
             index.delete(external_id=i, timestamp=i)
             index.update(
                 vector=data[i].astype(dtype),
@@ -660,6 +676,10 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
             )
             updated_ids[i] = i + update_ids_offset
 
+        ingestion_timestamps, base_sizes = load_metadata(index_uri)
+        assert ingestion_timestamps == [1]
+        assert base_sizes == [1000]
+
         index = index_class(uri=index_uri)
         _, result = index.query(queries, k=k, nprobe=partitions)
         assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
@@ -717,6 +737,11 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
 
         # Consolidate updates
         index = index.consolidate_updates()
+
+        ingestion_timestamps, base_sizes = load_metadata(index_uri)
+        assert ingestion_timestamps == [1, timestamp_end]
+        assert base_sizes == [1000, 1000]
+
         index = index_class(uri=index_uri)
         _, result = index.query(queries, k=k, nprobe=partitions)
         assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0

src/include/index/index_group.h

@@ -370,10 +370,13 @@ 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);
   }
-  auto get_all_ingestion_timestamps() {
+  auto get_all_ingestion_timestamps() const {
     return metadata_.ingestion_timestamps_;
   }
 
@@ -389,7 +392,7 @@ class base_index_group {
   auto append_base_size(size_t size) {
     metadata_.base_sizes_.push_back(size);
   }
-  auto get_all_base_sizes() {
+  auto get_all_base_sizes() const {
     return metadata_.base_sizes_;
   }
 
@@ -407,6 +410,10 @@ class base_index_group {
     metadata_.dimension_ = dim;
   }
 
+  auto get_timetravel_index() const {
+    return timetravel_index_;
+  }
+
   /**************************************************************************
    * Getters for names and uris
    **************************************************************************/

src/include/index/vamana_index.h

@@ -949,6 +949,13 @@ class vamana_index {
     return true;
   }
 
+  const vamana_index_group<vamana_index>& group() const {
+    if (!group_) {
+      throw std::runtime_error("No group available");
+    }
+    return *group_;
+  }
+
   /**
    * @brief Log statistics about the index
    */
@@ -1035,6 +1042,20 @@ class vamana_index {
                 << std::endl;
       return false;
     }
+    if (temporal_policy_.timestamp_start() !=
+        rhs.temporal_policy_.timestamp_start()) {
+      std::cout << "temporal_policy_.timestamp_start() != "
+                   "rhs.temporal_policy_.timestamp_start()"
+                << medoid_ << " ! = " << rhs.medoid_ << std::endl;
+      return false;
+    }
+    if (temporal_policy_.timestamp_end() !=
+        rhs.temporal_policy_.timestamp_end()) {
+      std::cout << "temporal_policy_.timestamp_end() != "
+                   "rhs.temporal_policy_.timestamp_end()"
+                << medoid_ << " ! = " << rhs.medoid_ << std::endl;
+      return false;
+    }
 
     return true;
   }