Support Vamana in the ObjectIndex (#366)

Author: jparismorgan

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

@@ -19,6 +19,7 @@
 from .module import validate_top_k
 from .storage_formats import STORAGE_VERSION
 from .storage_formats import storage_formats
+from .vamana_index import VamanaIndex
 
 try:
     from tiledb.vector_search.version import version as __version__

apis/python/src/tiledb/vector_search/object_api/object_index.py

@@ -11,8 +11,10 @@
 from tiledb.cloud.dag import Mode
 from tiledb.vector_search import FlatIndex
 from tiledb.vector_search import IVFFlatIndex
+from tiledb.vector_search import VamanaIndex
 from tiledb.vector_search import flat_index
 from tiledb.vector_search import ivf_flat_index
+from tiledb.vector_search import vamana_index
 from tiledb.vector_search.embeddings import ObjectEmbedding
 from tiledb.vector_search.object_readers import ObjectReader
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
@@ -53,6 +55,12 @@ def __init__(
                 self.index = IVFFlatIndex(
                     uri=uri, config=config, timestamp=timestamp, **kwargs
                 )
+            elif self.index_type == "VAMANA":
+                self.index = VamanaIndex(
+                    uri=uri, config=config, timestamp=timestamp, **kwargs
+                )
+            else:
+                raise ValueError(f"Unsupported index type {self.index_type}")
 
             self.object_reader_source_code = self.index.group.meta[
                 "object_reader_source_code"
@@ -428,6 +436,16 @@ def create(
                 config=config,
                 storage_version=storage_version,
             )
+        elif index_type == "VAMANA":
+            index = vamana_index.create(
+                uri=uri,
+                dimensions=dimensions,
+                vector_type=vector_type,
+                config=config,
+                storage_version=storage_version,
+            )
+        else:
+            raise ValueError(f"Unsupported index type {index_type}")
 
         group = tiledb.Group(uri, "w")
         group.meta["object_reader_source_code"] = get_source_code(object_reader)

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

@@ -92,7 +92,8 @@ def query_internal(
 
         assert queries.dtype == np.float32
         if opt_l < k:
-            raise ValueError(f"opt_l ({opt_l}) should be >= k ({k})")
+            warnings.warn(f"opt_l ({opt_l}) should be >= k ({k}), setting to k")
+            opt_l = k
 
         if queries.ndim == 1:
             queries = np.array([queries])

apis/python/test/test_index.py

@@ -279,10 +279,6 @@ def test_vamana_index(tmp_path):
         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)
-
     # Test that we can query with multiple query vectors.
     for i in range(5):
         query_and_check_distances(

apis/python/test/test_object_index.py

@@ -1,3 +1,4 @@
+import os
 from typing import Dict, List, OrderedDict, Tuple
 
 import numpy as np
@@ -10,6 +11,7 @@
 from tiledb.vector_search.object_readers import ObjectReader
 
 EMBED_DIM = 4
+INDEXES = ["FLAT", "IVF_FLAT", "VAMANA"]
 
 
 # TestEmbedding with vectors of EMBED_DIM size with all values being the id of the object
@@ -233,73 +235,73 @@ def df_filter(row):
     )
 
 
-def test_object_index_ivf_flat(tmp_path):
-    reader = TestReader(
-        object_id_start=0,
-        object_id_end=1000,
-        vector_dim_offset=0,
-    )
-    embedding = TestEmbedding()
-
-    index_uri = f"{tmp_path}/index"
+def test_object_index(tmp_path):
+    for index_type in INDEXES:
+        index_uri = os.path.join(tmp_path, f"object_index_{index_type}")
+        reader = TestReader(
+            object_id_start=0,
+            object_id_end=1000,
+            vector_dim_offset=0,
+        )
+        embedding = TestEmbedding()
 
-    index = object_index.create(
-        uri=index_uri,
-        index_type="IVF_FLAT",
-        object_reader=reader,
-        embedding=embedding,
-    )
+        index = object_index.create(
+            uri=index_uri,
+            index_type=index_type,
+            object_reader=reader,
+            embedding=embedding,
+        )
 
-    # Check initial ingestion
-    index.update_index(partitions=10)
-    evaluate_query(
-        index_uri=index_uri,
-        query_kwargs={"nprobe": 10},
-        dim_id=42,
-        vector_dim_offset=0,
-    )
+        # Check initial ingestion
+        index.update_index(partitions=10)
+        evaluate_query(
+            index_uri=index_uri,
+            query_kwargs={"nprobe": 10, "opt_l": 250},
+            dim_id=42,
+            vector_dim_offset=0,
+        )
 
-    # Check that updating the same data doesn't create duplicates
-    index = object_index.ObjectIndex(uri=index_uri)
-    index.update_index(partitions=10)
-    evaluate_query(
-        index_uri=index_uri,
-        query_kwargs={"nprobe": 10},
-        dim_id=42,
-        vector_dim_offset=0,
-    )
+        # Check that updating the same data doesn't create duplicates
+        index = object_index.ObjectIndex(uri=index_uri)
+        index.update_index(partitions=10)
+        evaluate_query(
+            index_uri=index_uri,
+            query_kwargs={"nprobe": 10, "opt_l": 500},
+            dim_id=42,
+            vector_dim_offset=0,
+        )
 
-    # Add new data with a new reader
-    reader = TestReader(
-        object_id_start=1000,
-        object_id_end=2000,
-        vector_dim_offset=0,
-    )
-    index = object_index.ObjectIndex(uri=index_uri)
-    index.update_object_reader(reader)
-    index.update_index(partitions=10)
-    evaluate_query(
-        index_uri=index_uri,
-        query_kwargs={"nprobe": 10},
-        dim_id=1042,
-        vector_dim_offset=0,
-    )
+        # Add new data with a new reader
+        reader = TestReader(
+            object_id_start=1000,
+            object_id_end=2000,
+            vector_dim_offset=0,
+        )
+        index = object_index.ObjectIndex(uri=index_uri)
+        index.update_object_reader(reader)
+        index.update_index(partitions=10)
+        evaluate_query(
+            index_uri=index_uri,
+            query_kwargs={"nprobe": 10, "opt_l": 500},
+            dim_id=1042,
+            vector_dim_offset=0,
+        )
 
-    # Check overwritting existing data
-    reader = TestReader(
-        object_id_start=1000,
-        object_id_end=2000,
-        vector_dim_offset=1000,
-    )
-    index = object_index.ObjectIndex(uri=index_uri)
-    index.update_object_reader(reader)
-    index.update_index(partitions=10)
-    evaluate_query(
-        index_uri=index_uri,
-        query_kwargs={"nprobe": 10},
-        dim_id=2042,
-        vector_dim_offset=1000,
-    )
+        # Check overwritting existing data
+        reader = TestReader(
+            object_id_start=1000,
+            object_id_end=2000,
+            vector_dim_offset=1000,
+        )
+        index = object_index.ObjectIndex(uri=index_uri)
+        index.update_object_reader(reader)
+        index.update_index(partitions=10)
+        evaluate_query(
+            index_uri=index_uri,
+            query_kwargs={"nprobe": 10, "opt_l": 500},
+            dim_id=2042,
+            vector_dim_offset=1000,
+        )
 
 
 def test_object_index_ivf_flat_cloud(tmp_path):