Re-enable IVF PQ tests (#450)

Author: jparismorgan

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

@@ -60,7 +60,6 @@ def __init__(
         )
         # TODO(SC-48710): Add support for `open_for_remote_query_execution`. We don't leave `self.index`` as `None` because we need to be able to call index.dimensions().
         self.index = vspy.IndexIVFPQ(self.ctx, uri, to_temporal_policy(timestamp))
-        # TODO(paris): This is incorrect - should be fixed when we fix consolidation.
         self.db_uri = self.group[
             storage_formats[self.storage_version]["PARTS_ARRAY_NAME"]
         ].uri

apis/python/test/test_index.py

@@ -350,7 +350,46 @@ def test_ivf_pq_index(tmp_path):
         index, np.array([[2, 2, 2]], dtype=np.float32), 2, [[0, 3]], [[2, 1]]
     )
 
-    # TODO(paris): Add tests for consolidation once we enable it.
+    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
+    )
+
+    # Test that we can query with multiple query vectors.
+    for i in range(5):
+        query_and_check_distances(
+            index,
+            np.array([[i, i, i], [i, i, i]], dtype=np.float32),
+            1,
+            [[0], [0]],
+            [[i], [i]],
+        )
+
+    # Test that we can query with k > 1.
+    query_and_check_distances(
+        index, np.array([[0, 0, 0]], dtype=np.float32), 2, [[0, 3]], [[0, 1]]
+    )
+
+    # Test that we can query with multiple query vectors and k > 1.
+    query_and_check_distances(
+        index,
+        np.array([[0, 0, 0], [4, 4, 4]], dtype=np.float32),
+        2,
+        [[0, 3], [0, 3]],
+        [[0, 1], [4, 3]],
+    )
+
+    vfs = tiledb.VFS()
+    assert vfs.dir_size(uri) > 0
+    Index.delete_index(uri=uri, config={})
+    assert vfs.dir_size(uri) == 0
 
 
 def test_delete_invalid_index(tmp_path):

apis/python/test/test_ingestion.py

@@ -824,12 +824,7 @@ def test_ingestion_with_updates(tmp_path):
         ingestion_timestamp = ingestion_timestamps[0]
 
         _, result = index.query(queries, k=k, nprobe=nprobe)
-        if index_type == "IVF_PQ":
-            # TODO(paris): We get 0.989 accuracy instead of 1.0. Investigate why - it should be 1.0
-            # when we have `nprobe = partitions` and `num_subspaces = dimensions`.
-            assert accuracy(result, gt_i) > 0.9
-            continue
-        assert accuracy(result, gt_i) == 1.0
+        assert accuracy(result, gt_i) >= (0.998 if index_type == "IVF_PQ" else 1.0)
 
         update_ids_offset = MAX_UINT64 - size
         updated_ids = {}
@@ -841,16 +836,22 @@ def test_ingestion_with_updates(tmp_path):
             updated_ids[i] = i + update_ids_offset
 
         _, result = index.query(queries, k=k, nprobe=nprobe)
-        assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+        assert accuracy(result, gt_i, updated_ids=updated_ids) >= (
+            0.998 if index_type == "IVF_PQ" else 1.0
+        )
 
         index = index.consolidate_updates(retrain_index=True, partitions=20)
         _, result = index.query(queries, k=k, nprobe=20)
-        assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+        assert accuracy(result, gt_i, updated_ids=updated_ids) >= (
+            0.998 if index_type == "IVF_PQ" else 1.0
+        )
 
         index_uri = move_local_index_to_new_location(index_uri)
         index = index_class(uri=index_uri)
         _, result = index.query(queries, k=k, nprobe=20)
-        assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+        assert accuracy(result, gt_i, updated_ids=updated_ids) >= (
+            0.998 if index_type == "IVF_PQ" else 1.0
+        )
 
         ingestion_timestamps, base_sizes = load_metadata(index_uri)
         assert base_sizes == [1000, 1000]