Fix type-erased indexes writing fragments at timestamp=0, thus fixing IVF PQ time travel (#425)

Author: jparismorgan

apis/python/test/test_ingestion.py

@@ -957,10 +957,6 @@ def test_ingestion_with_updates_and_timetravel(tmp_path):
         assert ingestion_timestamps == [1]
         assert base_sizes == [size]
 
-        if index_type == "IVF_PQ":
-            # TODO(SC-48897): Fix time travelling for IVF_PQ and re-enable.
-            continue
-
         if index_type == "IVF_FLAT":
             assert index.partitions == partitions
 
@@ -1226,9 +1222,6 @@ def test_ingestion_with_additions_and_timetravel(tmp_path):
         )
         if index_type == "IVF_FLAT":
             assert index.partitions == partitions
-        if index_type == "IVF_PQ":
-            # TODO(SC-48897): Fix time travelling for IVF_PQ and re-enable.
-            continue
         _, result = index.query(queries, k=k, nprobe=partitions)
         assert accuracy(result, gt_i) == 1.0
 

src/include/index/ivf_pq_index.h

@@ -1040,6 +1040,15 @@ class ivf_pq_index {
       write_group.append_num_partitions(num_partitions_);
     }
 
+    // When creating from Python we initially call write_index() at timestamp 0.
+    // The goal here is just to create the arrays and save metadata. Return here
+    // so that we don't write the arrays, as if we write with timestamp=0 then
+    // TileDB Core will interpret this as the current timestamp instead, leading
+    // to array fragments created at the current time.
+    if (temporal_policy_.timestamp_end() == 0) {
+      return true;
+    }
+
     // flat_ivf_centroids_, cluster_centroids_, distance_tables_
     // pq_ivf_centroids_, partitioned_pq_vectors_, unpartitioned_pq_vectors_
 

src/include/index/vamana_index.h

@@ -757,6 +757,15 @@ class vamana_index {
       write_group.append_num_edges(graph_.num_edges());
     }
 
+    // When creating from Python we initially call write_index() at timestamp 0.
+    // The goal here is just to create the arrays and save metadata. Return here
+    // so that we don't write the arrays, as if we write with timestamp=0 then
+    // TileDB Core will interpret this as the current timestamp instead, leading
+    // to array fragments created at the current time.
+    if (temporal_policy_.timestamp_end() == 0) {
+      return true;
+    }
+
     write_matrix(
         ctx,
         feature_vectors_,

src/include/test/unit_ivf_pq_index.cc

@@ -695,3 +695,69 @@ TEST_CASE("query simple", "[ivf_pq_index]") {
     }
   }
 }
+
+TEST_CASE("ivf_pq_index query index written twice", "[ivf_pq_index]") {
+  tiledb::Context ctx;
+  tiledb::VFS vfs(ctx);
+  std::string index_uri =
+      (std::filesystem::temp_directory_path() / "tmp_ivf_pq_index").string();
+  std::cout << "index_uri: " << index_uri << std::endl;
+  if (vfs.is_dir(index_uri)) {
+    vfs.remove_dir(index_uri);
+  }
+
+  using feature_type_type = uint8_t;
+  using id_type_type = uint32_t;
+  using partitioning_index_type_type = uint32_t;
+  auto feature_type = "uint8";
+  auto id_type = "uint32";
+  auto partitioning_index_type = "uint32";
+  size_t dimensions = 3;
+  size_t n_list = 1;
+  size_t num_subspaces = 1;
+  float convergence_tolerance = 0.00003f;
+  size_t max_iterations = 3;
+
+  // Write the empty index.
+  {
+    auto index = ivf_pq_index<
+        feature_type_type,
+        id_type_type,
+        partitioning_index_type_type>(n_list, dimensions / 2);
+    auto data =
+        ColMajorMatrixWithIds<feature_type_type, id_type_type>(dimensions, 0);
+    index.train(data, data.raveled_ids());
+    index.add(data, data.raveled_ids());
+    index.write_index(ctx, index_uri, TemporalPolicy(TimeTravel, 0));
+  }
+
+  // Train the index at timestamp 99.
+  {
+    auto index = ivf_pq_index<
+        feature_type_type,
+        id_type_type,
+        partitioning_index_type_type>(ctx, index_uri);
+    auto data = ColMajorMatrixWithIds<feature_type_type, id_type_type>{
+        {{1, 1, 1}, {2, 2, 2}, {3, 3, 3}, {4, 4, 4}}, {1, 2, 3, 4}};
+    index.train(data, data.raveled_ids());
+    index.add(data, data.raveled_ids());
+    index.write_index(ctx, index_uri, TemporalPolicy(TimeTravel, 99));
+  }
+
+  // Load the index and query.
+  {
+    auto index = ivf_pq_index<
+        feature_type_type,
+        id_type_type,
+        partitioning_index_type_type>(ctx, index_uri);
+    auto queries = ColMajorMatrix<feature_type_type>{
+        {1, 1, 1}, {2, 2, 2}, {3, 3, 3}, {4, 4, 4}};
+    auto&& [scores, ids] = index.query_infinite_ram(queries, 1, n_list);
+    CHECK(std::equal(
+        scores.data(),
+        scores.data() + 4,
+        std::vector<float>{0, 0, 0, 0}.begin()));
+    CHECK(std::equal(
+        ids.data(), ids.data() + 4, std::vector<uint32_t>{1, 2, 3, 4}.begin()));
+  }
+}