Exposes all versions of query functions to Python (#79)

This exposes all versions of query functions to Python. 
Also did some minor tweeks for ingestion and ran black to format the code.

Author: NikolaosPapailiou

apis/python/setup.py

@@ -4,7 +4,6 @@
 
 
 def get_cmake_overrides():
-
     conf = list()
 
     tiledb_dir = os.environ.get("TILEDB_DIR", None)

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

@@ -2,7 +2,16 @@
 from .ingestion import ingest
 from .module import load_as_array
 from .module import load_as_matrix
-from .module import query_vq, query_kmeans, validate_top_k, array_to_matrix, ivf_index, ivf_index_tdb, partition_ivf_index
+from .module import (
+    query_vq,
+    ivf_query,
+    ivf_query_ram,
+    validate_top_k,
+    array_to_matrix,
+    ivf_index,
+    ivf_index_tdb,
+    partition_ivf_index,
+)
 
 __all__ = [
     "FlatIndex",
@@ -11,10 +20,11 @@
     "load_as_matrix",
     "ingest",
     "query_vq",
-    "query_kmeans",
+    "ivf_query",
+    "ivf_query_ram",
     "validate_top_k",
     "ivf_index",
     "ivf_index_tdb",
     "array_to_matrix",
-    "partition_ivf_index"
+    "partition_ivf_index",
 ]

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

@@ -76,22 +76,40 @@ class IVFFlatIndex(Index):
         URI of datataset
     dtype: numpy.dtype
         datatype float32 or uint8
+    memory_budget: int
+        Main memory budget. If not provided no memory budget is applied.
     """
 
-    def __init__(self, uri, dtype: np.dtype):
+    def __init__(
+        self, uri, dtype: np.dtype, memory_budget: int = -1, ctx: "Ctx" = None
+    ):
         self.parts_db_uri = os.path.join(uri, "parts.tdb")
         self.centroids_uri = os.path.join(uri, "centroids.tdb")
         self.index_uri = os.path.join(uri, "index.tdb")
         self.ids_uri = os.path.join(uri, "ids.tdb")
         self.dtype = dtype
+        self.memory_budget = memory_budget
+        self.ctx = ctx
+        if ctx is None:
+            self.ctx = Ctx({})
+
+        # TODO pass in a context
+        if self.memory_budget == -1:
+            self._db = load_as_matrix(self.parts_db_uri)
+            self._ids = read_vector_u64(self.ctx, self.ids_uri)
 
-        ctx = Ctx({})  # TODO pass in a context
-        self._db = load_as_matrix(self.parts_db_uri)
         self._centroids = load_as_matrix(self.centroids_uri)
-        self._index = read_vector_u64(ctx, self.index_uri)
-        self._ids = read_vector_u64(ctx, self.ids_uri)
+        self._index = read_vector_u64(self.ctx, self.index_uri)
 
-    def query(self, targets: np.ndarray, k=10, nqueries=10, nthreads=8, nprobe=1):
+    def query(
+        self,
+        targets: np.ndarray,
+        k=10,
+        nqueries=10,
+        nthreads=8,
+        nprobe=1,
+        use_nuv_implementation: bool = False,
+    ):
         """
         Open a flat index
 
@@ -107,21 +125,42 @@ def query(self, targets: np.ndarray, k=10, nqueries=10, nthreads=8, nprobe=1):
             Number of threads to use for queyr
         nprobe: int
             number of probes
+        use_nuv_implementation: bool
+            wether to use the nuv query implementation. Default: False
         """
         assert targets.dtype == np.float32
 
         targets_m = array_to_matrix(targets)
+        if self.memory_budget == -1:
+            r = ivf_query_ram(
+                self.dtype,
+                self._db,
+                self._centroids,
+                targets_m,
+                self._index,
+                self._ids,
+                nprobe=nprobe,
+                k_nn=k,
+                nth=True,  # ??
+                nthreads=nthreads,
+                ctx=self.ctx,
+                use_nuv_implementation=use_nuv_implementation,
+            )
+        else:
+            r = ivf_query(
+                self.dtype,
+                self.parts_db_uri,
+                self._centroids,
+                targets_m,
+                self._index,
+                self.ids_uri,
+                nprobe=nprobe,
+                k_nn=k,
+                memory_budget=self.memory_budget,
+                nth=True,  # ??
+                nthreads=nthreads,
+                ctx=self.ctx,
+                use_nuv_implementation=use_nuv_implementation,
+            )
 
-        r = query_kmeans(
-            self._db.dtype,
-            self._db,
-            self._centroids,
-            targets_m,
-            self._index,
-            self._ids,
-            nprobe=nprobe,
-            k_nn=k,
-            nth=True,  # ??
-            nthreads=nthreads,
-        )
         return np.array(r)