Add open method for Index class (#503)

This adds `open` method for `Index` class which can be used to open an index without knowing its `index_type`

Author: NikolaosPapailiou

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

@@ -4,6 +4,7 @@
 from . import utils
 from .flat_index import FlatIndex
 from .index import Index
+from .index import open
 from .ingestion import ingest
 from .ivf_flat_index import IVFFlatIndex
 from .ivf_pq_index import IVFPQIndex
@@ -34,6 +35,7 @@
     "VamanaIndex",
     "IVFPQIndex",
     "Mode",
+    "open",
     "load_as_array",
     "load_as_matrix",
     "ingest",

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

@@ -46,6 +46,7 @@ def __init__(
         config: Optional[Mapping[str, Any]] = None,
         timestamp=None,
         open_for_remote_query_execution: bool = False,
+        group: tiledb.Group = None,
         **kwargs,
     ):
         self.index_open_kwargs = {
@@ -60,6 +61,7 @@ def __init__(
             config=config,
             timestamp=timestamp,
             open_for_remote_query_execution=open_for_remote_query_execution,
+            group=group,
         )
         self._index = None
         self.db_uri = self.group[

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

@@ -2,6 +2,8 @@
 import json
 import os
 import time
+from abc import ABCMeta
+from abc import abstractmethod
 from typing import Any, Mapping, Optional
 
 from tiledb.cloud.dag import Mode
@@ -16,7 +18,7 @@
 DATASET_TYPE = "vector_search"
 
 
-class Index:
+class Index(metaclass=ABCMeta):
     """
     Abstract Vector Index class.
 
@@ -42,22 +44,26 @@ class Index:
         If `False`, load index data in main memory locally. Note that you can still use a taskgraph for query execution, you'll just end up loading the data both on your local machine and in the cloud taskgraph.
     """
 
+    @abstractmethod
     def __init__(
         self,
         uri: str,
-        open_for_remote_query_execution: bool,
+        open_for_remote_query_execution: bool = False,
         config: Optional[Mapping[str, Any]] = None,
         timestamp=None,
+        group: tiledb.Group = None,
     ):
         # If the user passes a tiledb python Config object convert to a dictionary
         if isinstance(config, tiledb.Config):
             config = dict(config)
-
         self.uri = uri
         self.open_for_remote_query_execution = open_for_remote_query_execution
         self.config = config
         self.ctx = vspy.Ctx(config)
-        self.group = tiledb.Group(self.uri, "r", ctx=tiledb.Ctx(config))
+        if group is not None:
+            self.group = group
+        else:
+            self.group = tiledb.Group(self.uri, "r", ctx=tiledb.Ctx(config))
         self.storage_version = self.group.meta.get("storage_version", "0.1")
         try:
             self.distance_metric = vspy.DistanceMetric(
@@ -688,6 +694,7 @@ def clear_history(
                 raise ValueError(f"Unsupported index_type: {index_type}")
             group.close()
 
+    @abstractmethod
     def get_dimensions(self):
         """
         Abstract method implemented by all Vector Index implementations.
@@ -696,6 +703,7 @@ def get_dimensions(self):
         """
         raise NotImplementedError
 
+    @abstractmethod
     def query_internal(self, queries: np.ndarray, k: int, **kwargs):
         """
         Abstract method implemented by all Vector Index implementations.
@@ -868,3 +876,80 @@ def create_metadata(
         group.meta["has_updates"] = False
         group.meta["distance_metric"] = int(distance_metric)
         group.close()
+
+
+"""
+Factory method that opens a vector index.
+
+Retrieves the `index_type` from the index group metadata and instantiates the appropriate `Index` subclass.
+
+Parameters
+----------
+uri: str
+    URI of the index.
+config: Optional[Mapping[str, Any]]
+    TileDB config dictionary.
+timestamp: int or tuple(int)
+    If int, open the index at a given timestamp.
+    If tuple, open at the given start and end timestamps.
+open_for_remote_query_execution: bool
+    If `True`, do not load any index data in main memory locally, and instead load index data in the TileDB Cloud taskgraph created when a non-`None` `driver_mode` is passed to `query()`.
+    If `False`, load index data in main memory locally. Note that you can still use a taskgraph for query execution, you'll just end up loading the data both on your local machine and in the cloud taskgraph.
+kwargs:
+    Additional arguments to be passed to the `Index` subclass constructor.
+"""
+
+
+def open(
+    uri: str,
+    open_for_remote_query_execution: bool = False,
+    config: Optional[Mapping[str, Any]] = None,
+    timestamp=None,
+    **kwargs,
+) -> Index:
+    from tiledb.vector_search.flat_index import FlatIndex
+    from tiledb.vector_search.ivf_flat_index import IVFFlatIndex
+    from tiledb.vector_search.ivf_pq_index import IVFPQIndex
+    from tiledb.vector_search.vamana_index import VamanaIndex
+
+    group = tiledb.Group(uri, "r")
+    index_type = group.meta["index_type"]
+    if index_type == "FLAT":
+        return FlatIndex(
+            uri=uri,
+            open_for_remote_query_execution=open_for_remote_query_execution,
+            config=config,
+            timestamp=timestamp,
+            group=group,
+            **kwargs,
+        )
+    elif index_type == "IVF_FLAT":
+        return IVFFlatIndex(
+            uri=uri,
+            open_for_remote_query_execution=open_for_remote_query_execution,
+            config=config,
+            timestamp=timestamp,
+            group=group,
+            **kwargs,
+        )
+    elif index_type == "VAMANA":
+        return VamanaIndex(
+            uri=uri,
+            open_for_remote_query_execution=open_for_remote_query_execution,
+            config=config,
+            timestamp=timestamp,
+            group=group,
+            **kwargs,
+        )
+    elif index_type == "IVF_PQ":
+        return IVFPQIndex(
+            uri=uri,
+            open_for_remote_query_execution=open_for_remote_query_execution,
+            config=config,
+            timestamp=timestamp,
+            group=group,
+            **kwargs,
+        )
+    else:
+        group.close()
+        raise ValueError(f"Unsupported index type {index_type}")

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

@@ -218,6 +218,7 @@ def ingest(
     from tiledb.cloud.utilities import get_logger
     from tiledb.cloud.utilities import set_aws_context
     from tiledb.vector_search import flat_index
+    from tiledb.vector_search import index
     from tiledb.vector_search import ivf_flat_index
     from tiledb.vector_search import ivf_pq_index
     from tiledb.vector_search import vamana_index
@@ -3061,16 +3062,4 @@ def consolidate_and_vacuum(
             group.close()
 
         consolidate_and_vacuum(index_group_uri=index_group_uri, config=config)
-
-        if index_type == "FLAT":
-            return flat_index.FlatIndex(uri=index_group_uri, config=config)
-        elif index_type == "VAMANA":
-            return vamana_index.VamanaIndex(uri=index_group_uri, config=config)
-        elif index_type == "IVF_FLAT":
-            return ivf_flat_index.IVFFlatIndex(
-                uri=index_group_uri, memory_budget=1000000, config=config
-            )
-        elif index_type == "IVF_PQ":
-            return ivf_pq_index.IVFPQIndex(uri=index_group_uri, config=config)
-        else:
-            raise ValueError(f"Not supported index_type {index_type}")
+        return index.open(uri=index_group_uri, memory_budget=1000000, config=config)

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

@@ -84,6 +84,7 @@ def __init__(
         timestamp=None,
         memory_budget: int = -1,
         open_for_remote_query_execution: bool = False,
+        group: tiledb.Group = None,
         **kwargs,
     ):
         self.index_open_kwargs = {
@@ -99,6 +100,7 @@ def __init__(
             config=config,
             timestamp=timestamp,
             open_for_remote_query_execution=open_for_remote_query_execution,
+            group=group,
         )
         self.db_uri = self.group[
             storage_formats[self.storage_version]["PARTS_ARRAY_NAME"]

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

@@ -49,6 +49,7 @@ def __init__(
         timestamp=None,
         memory_budget: int = -1,
         open_for_remote_query_execution: bool = False,
+        group: tiledb.Group = None,
         **kwargs,
     ):
         self.index_open_kwargs = {
@@ -64,6 +65,7 @@ def __init__(
             config=config,
             timestamp=timestamp,
             open_for_remote_query_execution=open_for_remote_query_execution,
+            group=group,
         )
         # 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))

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

@@ -53,6 +53,7 @@ def __init__(
         config: Optional[Mapping[str, Any]] = None,
         timestamp=None,
         open_for_remote_query_execution: bool = False,
+        group: tiledb.Group = None,
         **kwargs,
     ):
         self.index_open_kwargs = {
@@ -67,6 +68,7 @@ def __init__(
             config=config,
             timestamp=timestamp,
             open_for_remote_query_execution=open_for_remote_query_execution,
+            group=group,
         )
         # 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.IndexVamana(self.ctx, uri, to_temporal_policy(timestamp))

apis/python/test/test_index.py

@@ -12,13 +12,11 @@
 from tiledb.vector_search import flat_index
 from tiledb.vector_search import ivf_flat_index
 from tiledb.vector_search import ivf_pq_index
+from tiledb.vector_search import open
 from tiledb.vector_search import vamana_index
-from tiledb.vector_search.flat_index import FlatIndex
 from tiledb.vector_search.index import DATASET_TYPE
 from tiledb.vector_search.index import create_metadata
 from tiledb.vector_search.ingestion import ingest
-from tiledb.vector_search.ivf_flat_index import IVFFlatIndex
-from tiledb.vector_search.ivf_pq_index import IVFPQIndex
 from tiledb.vector_search.utils import MAX_FLOAT32
 from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import is_type_erased_index
@@ -423,9 +421,8 @@ def test_delete_index(tmp_path):
     vfs = tiledb.VFS()
 
     indexes = ["FLAT", "IVF_FLAT", "VAMANA", "IVF_PQ"]
-    index_classes = [FlatIndex, IVFFlatIndex, VamanaIndex, IVFPQIndex]
     data = np.array([[1.0, 1.1, 1.2, 1.3], [2.0, 2.1, 2.2, 2.3]], dtype=np.float32)
-    for index_type, index_class in zip(indexes, index_classes):
+    for index_type in indexes:
         index_uri = os.path.join(tmp_path, f"array_{index_type}")
         ingest(
             index_type=index_type,
@@ -436,7 +433,7 @@ def test_delete_index(tmp_path):
         Index.delete_index(uri=index_uri, config={})
         assert vfs.dir_size(index_uri) == 0
         with pytest.raises(tiledb.TileDBError) as error:
-            index_class(uri=index_uri)
+            open(uri=index_uri)
         assert "does not exist" in str(error.value)
 
 

apis/python/test/test_ingestion.py

@@ -8,10 +8,10 @@
 
 from tiledb.cloud.dag import Mode
 from tiledb.vector_search import _tiledbvspy as vspy
+from tiledb.vector_search import open
 from tiledb.vector_search.index import Index
 from tiledb.vector_search.ingestion import TrainingSamplingPolicy
 from tiledb.vector_search.ingestion import ingest
-from tiledb.vector_search.ivf_flat_index import IVFFlatIndex
 from tiledb.vector_search.module import array_to_matrix
 from tiledb.vector_search.module import kmeans_fit
 from tiledb.vector_search.module import kmeans_predict
@@ -72,7 +72,7 @@ def test_vamana_ingestion_u8(tmp_path):
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index_ram = VamanaIndex(uri=index_uri)
+    index_ram = open(uri=index_uri)
     _, result = index_ram.query(queries, k=k)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
@@ -100,7 +100,7 @@ def test_flat_ingestion_u8(tmp_path):
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index_ram = FlatIndex(uri=index_uri)
+    index_ram = open(uri=index_uri)
     _, result = index_ram.query(queries, k=k)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
@@ -124,7 +124,7 @@ def test_flat_ingestion_f32(tmp_path):
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index_ram = FlatIndex(uri=index_uri)
+    index_ram = open(uri=index_uri)
     _, result = index_ram.query(queries, k=k)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
@@ -157,7 +157,7 @@ def test_flat_ingestion_external_id_u8(tmp_path):
     )
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index_ram = FlatIndex(uri=index_uri)
+    index_ram = open(uri=index_uri)
     _, result = index_ram.query(queries, k=k)
     assert (
         accuracy(result, gt_i, external_ids_offset=external_ids_offset)
@@ -190,7 +190,7 @@ def test_ivf_flat_ingestion_u8(tmp_path):
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index_ram = IVFFlatIndex(uri=index_uri, memory_budget=int(size / 10))
+    index_ram = open(uri=index_uri, memory_budget=int(size / 10))
     _, result = index_ram.query(queries, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
@@ -237,7 +237,7 @@ def test_ivf_pq_ingestion_u8(tmp_path):
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index_ram = IVFPQIndex(uri=index_uri, memory_budget=int(size / 10))
+    index_ram = open(uri=index_uri, memory_budget=int(size / 10))
     _, result = index_ram.query(queries, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
@@ -1698,7 +1698,7 @@ def test_ivf_flat_ingestion_with_training_source_uri_f32(tmp_path):
     )
 
     index_uri = move_local_index_to_new_location(index_uri)
-    index = IVFFlatIndex(uri=index_uri)
+    index = open(uri=index_uri)
     query_and_check_equals(
         index=index, queries=queries, expected_result_d=[[0]], expected_result_i=[[1]]
     )
@@ -1796,7 +1796,7 @@ def test_ivf_flat_ingestion_with_training_source_uri_tdb(tmp_path):
     index_uri = move_local_index_to_new_location(index_uri)
 
     # Load the index again and query.
-    index = IVFFlatIndex(uri=index_uri)
+    index = open(uri=index_uri)
 
     query_and_check_equals(
         index=index,
@@ -1823,7 +1823,7 @@ def test_ivf_flat_ingestion_with_training_source_uri_tdb(tmp_path):
     # Clear the index history, load, update, and query.
     Index.clear_history(uri=index_uri, timestamp=index.latest_ingestion_timestamp - 1)
 
-    index = IVFFlatIndex(uri=index_uri)
+    index = open(uri=index_uri)
 
     update_vectors = np.empty([2], dtype=object)
     update_vectors[0] = np.array([11.0, 11.1, 11.2, 11.3], dtype=np.dtype(np.float32))
@@ -1918,7 +1918,7 @@ def test_ivf_flat_ingestion_with_training_source_uri_numpy(tmp_path):
     # Test we can load the index again and query, update, and consolidate.
     ################################################################################################
     index_uri = move_local_index_to_new_location(index_uri)
-    index = IVFFlatIndex(uri=index_uri)
+    index = open(uri=index_uri)
 
     queries = np.array([data[1]], dtype=np.float32)
     query_and_check_equals(