Fix format and tests

Author: Nikos Papailiou

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

@@ -6,6 +6,7 @@
 import numpy as np
 from tiledb.vector_search.module import *
 from tiledb.cloud.dag import Mode
+from typing import Any, Mapping
 
 CENTROIDS_ARRAY_NAME = "centroids.tdb"
 INDEX_ARRAY_NAME = "index.tdb"
@@ -39,13 +40,17 @@ class FlatIndex(Index):
         Optional name of partitions
     """
 
-    def __init__(self, uri: str, dtype: np.dtype, parts_name: str = "parts.tdb", ctx: "Ctx" = None):
+    def __init__(
+        self,
+        uri: str,
+        dtype: np.dtype,
+        parts_name: str = "parts.tdb",
+        config: Optional[Mapping[str, Any]] = None,
+    ):
         self.uri = uri
         self.dtype = dtype
         self._index = None
-        self.ctx = ctx
-        if ctx is None:
-            self.ctx = Ctx({})
+        self.ctx = Ctx(config)
 
         self._db = load_as_matrix(os.path.join(uri, parts_name), ctx=self.ctx)
 
@@ -106,7 +111,11 @@ class IVFFlatIndex(Index):
     """
 
     def __init__(
-        self, uri, dtype: np.dtype, memory_budget: int = -1, ctx: "Ctx" = None
+        self,
+        uri,
+        dtype: np.dtype,
+        memory_budget: int = -1,
+        config: Optional[Mapping[str, Any]] = None,
     ):
         group = tiledb.Group(uri)
         self.parts_db_uri = group[PARTS_ARRAY_NAME].uri
@@ -115,9 +124,7 @@ def __init__(
         self.ids_uri = group[IDS_ARRAY_NAME].uri
         self.dtype = dtype
         self.memory_budget = memory_budget
-        self.ctx = ctx
-        if ctx is None:
-            self.ctx = Ctx({})
+        self.ctx = Ctx(config)
 
         # TODO pass in a context
         if self.memory_budget == -1:

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

@@ -2,8 +2,7 @@
 from functools import partial
 
 from tiledb.cloud.dag import Mode
-from tiledb.vector_search.index import FlatIndex
-from tiledb.vector_search.index import IVFFlatIndex
+from tiledb.vector_search.index import FlatIndex, IVFFlatIndex, Index
 
 
 def ingest(
@@ -23,7 +22,7 @@ def ingest(
     verbose: bool = False,
     trace_id: Optional[str] = None,
     mode: Mode = Mode.LOCAL,
-) -> FlatIndex:
+) -> Index:
     """
     Ingest vectors into TileDB.
 
@@ -271,9 +270,7 @@ def create_arrays(
                 )
                 logger.debug(centroids_schema)
                 tiledb.Array.create(centroids_uri, centroids_schema)
-                group.add(
-                    centroids_uri, name=CENTROIDS_ARRAY_NAME
-                )
+                group.add(centroids_uri, name=CENTROIDS_ARRAY_NAME)
 
             if not tiledb.array_exists(index_uri):
                 logger.debug("Creating index array")
@@ -354,7 +351,9 @@ def create_arrays(
             except tiledb.TileDBError as err:
                 message = str(err)
                 if "already exists" in message:
-                    logger.debug(f"Group '{partial_write_array_dir_uri}' already exists")
+                    logger.debug(
+                        f"Group '{partial_write_array_dir_uri}' already exists"
+                    )
                 raise err
             partial_write_array_group = tiledb.Group(partial_write_array_dir_uri, "w")
             group.add(partial_write_array_dir_uri, name=PARTIAL_WRITE_ARRAY_DIR)
@@ -364,10 +363,16 @@ def create_arrays(
             except tiledb.TileDBError as err:
                 message = str(err)
                 if "already exists" in message:
-                    logger.debug(f"Group '{partial_write_array_index_uri}' already exists")
+                    logger.debug(
+                        f"Group '{partial_write_array_index_uri}' already exists"
+                    )
                 raise err
-            partial_write_array_group.add(partial_write_array_index_uri, name=INDEX_ARRAY_NAME)
-            partial_write_array_index_group = tiledb.Group(partial_write_array_index_uri, "w")
+            partial_write_array_group.add(
+                partial_write_array_index_uri, name=INDEX_ARRAY_NAME
+            )
+            partial_write_array_index_group = tiledb.Group(
+                partial_write_array_index_uri, "w"
+            )
 
             if not tiledb.array_exists(partial_write_array_ids_uri):
                 logger.debug("Creating temp ids array")
@@ -389,7 +394,9 @@ def create_arrays(
                 )
                 logger.debug(ids_schema)
                 tiledb.Array.create(partial_write_array_ids_uri, ids_schema)
-                partial_write_array_group.add(partial_write_array_ids_uri, name=IDS_ARRAY_NAME)
+                partial_write_array_group.add(
+                    partial_write_array_ids_uri, name=IDS_ARRAY_NAME
+                )
 
             if not tiledb.array_exists(partial_write_array_parts_uri):
                 logger.debug("Creating temp parts array")
@@ -420,10 +427,12 @@ def create_arrays(
                 logger.debug(parts_schema)
                 logger.debug(partial_write_array_parts_uri)
                 tiledb.Array.create(partial_write_array_parts_uri, parts_schema)
-                partial_write_array_group.add(partial_write_array_parts_uri, name=PARTS_ARRAY_NAME)
+                partial_write_array_group.add(
+                    partial_write_array_parts_uri, name=PARTS_ARRAY_NAME
+                )
 
             for part in range(input_vectors_work_tasks):
-                part_index_uri = partial_write_array_index_uri+"/"+str(part)
+                part_index_uri = partial_write_array_index_uri + "/" + str(part)
                 if not tiledb.array_exists(part_index_uri):
                     logger.debug(f"Creating part array {part_index_uri}")
                     index_array_rows_dim = tiledb.Dim(
@@ -867,8 +876,12 @@ def ingest_vectors_udf(
         partial_write_array_group = tiledb.Group(partial_write_array_dir_uri)
         partial_write_array_ids_uri = partial_write_array_group[IDS_ARRAY_NAME].uri
         partial_write_array_parts_uri = partial_write_array_group[PARTS_ARRAY_NAME].uri
-        partial_write_array_index_dir_uri = partial_write_array_group[INDEX_ARRAY_NAME].uri
-        partial_write_array_index_group = tiledb.Group(partial_write_array_index_dir_uri)
+        partial_write_array_index_dir_uri = partial_write_array_group[
+            INDEX_ARRAY_NAME
+        ].uri
+        partial_write_array_index_group = tiledb.Group(
+            partial_write_array_index_dir_uri
+        )
 
         for part in range(start, end, batch):
             part_end = part + batch
@@ -877,7 +890,9 @@ def ingest_vectors_udf(
 
             part_name = str(part) + "-" + str(part_end)
 
-            partial_write_array_index_uri = partial_write_array_index_group[str(int(start / batch))].uri
+            partial_write_array_index_uri = partial_write_array_index_group[
+                str(int(start / batch))
+            ].uri
             logger.debug("Input vectors start_pos: %d, end_pos: %d", part, part_end)
             if source_type == "TILEDB_ARRAY":
                 logger.debug("Start indexing")
@@ -932,8 +947,12 @@ def compute_partition_indexes_udf(
             index_array_uri = group[INDEX_ARRAY_NAME].uri
             partial_write_array_dir_uri = group[PARTIAL_WRITE_ARRAY_DIR].uri
             partial_write_array_group = tiledb.Group(partial_write_array_dir_uri)
-            partial_write_array_index_dir_uri = partial_write_array_group[INDEX_ARRAY_NAME].uri
-            partial_write_array_index_group = tiledb.Group(partial_write_array_index_dir_uri)
+            partial_write_array_index_dir_uri = partial_write_array_group[
+                INDEX_ARRAY_NAME
+            ].uri
+            partial_write_array_index_group = tiledb.Group(
+                partial_write_array_index_dir_uri
+            )
             partition_sizes = np.zeros(partitions)
             indexes = np.zeros(partitions + 1).astype(np.uint64)
             for part in partial_write_array_index_group:
@@ -978,9 +997,15 @@ def consolidate_partition_udf(
             partial_write_array_dir_uri = group[PARTIAL_WRITE_ARRAY_DIR].uri
             partial_write_array_group = tiledb.Group(partial_write_array_dir_uri)
             partial_write_array_ids_uri = partial_write_array_group[IDS_ARRAY_NAME].uri
-            partial_write_array_parts_uri = partial_write_array_group[PARTS_ARRAY_NAME].uri
-            partial_write_array_index_dir_uri = partial_write_array_group[INDEX_ARRAY_NAME].uri
-            partial_write_array_index_group = tiledb.Group(partial_write_array_index_dir_uri)
+            partial_write_array_parts_uri = partial_write_array_group[
+                PARTS_ARRAY_NAME
+            ].uri
+            partial_write_array_index_dir_uri = partial_write_array_group[
+                INDEX_ARRAY_NAME
+            ].uri
+            partial_write_array_index_group = tiledb.Group(
+                partial_write_array_index_dir_uri
+            )
             index_array_uri = group[INDEX_ARRAY_NAME].uri
             ids_array_uri = group[IDS_ARRAY_NAME].uri
             parts_array_uri = group[PARTS_ARRAY_NAME].uri
@@ -1342,7 +1367,7 @@ def consolidate_and_vacuum(
             if vfs.is_dir(partial_write_array_dir_uri):
                 vfs.remove_dir(partial_write_array_dir_uri)
 
-    with tiledb.scope_ctx(ctx_or_config=config) as ctx:
+    with tiledb.scope_ctx(ctx_or_config=config):
         logger = setup(config, verbose)
         logger.debug("Ingesting Vectors into %r", array_uri)
         try:
@@ -1462,6 +1487,8 @@ def consolidate_and_vacuum(
         consolidate_and_vacuum(array_uri=array_uri, config=config)
 
         if index_type == "FLAT":
-            return FlatIndex(uri=array_uri, dtype=vector_type, ctx=ctx)
+            return FlatIndex(uri=array_uri, dtype=vector_type, config=config)
         elif index_type == "IVF_FLAT":
-            return IVFFlatIndex(uri=array_uri, dtype=vector_type, memory_budget=1000000, ctx=ctx)
+            return IVFFlatIndex(
+                uri=array_uri, dtype=vector_type, memory_budget=1000000, config=config
+            )

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

@@ -18,8 +18,8 @@ def load_as_matrix(path: str, nqueries: int = 0, ctx: "Ctx" = None):
         Array path
     nqueries: int
         Number of queries
-    config: Dict
-        TileDB configuration parameters
+    ctx: Ctx
+        TileDB context
     """
     if ctx is None:
         ctx = Ctx({})
@@ -44,7 +44,7 @@ def load_as_matrix(path: str, nqueries: int = 0, ctx: "Ctx" = None):
     return m
 
 
-def load_as_array(path, return_matrix: bool = False,  ctx: "Ctx" = None):
+def load_as_array(path, return_matrix: bool = False, ctx: "Ctx" = None):
     """
     Load array as array class
 

apis/python/test/common.py

@@ -96,7 +96,7 @@ def create_random_dataset_f32(nb, d, nq, k, path):
     from sklearn.datasets import make_blobs
     from sklearn.neighbors import NearestNeighbors
 
-    #print(f"Preparing datasets with {nb} random points and {nq} queries.")
+    # print(f"Preparing datasets with {nb} random points and {nq} queries.")
     os.mkdir(path)
     X, _ = make_blobs(n_samples=nb + nq, n_features=d, centers=nq, random_state=1)
 
@@ -111,7 +111,7 @@ def create_random_dataset_f32(nb, d, nq, k, path):
         np.array([nq, d], dtype="uint32").tofile(f)
         queries.astype("float32").tofile(f)
 
-    #print("Computing groundtruth")
+    # print("Computing groundtruth")
 
     nbrs = NearestNeighbors(n_neighbors=k, metric="euclidean", algorithm="brute").fit(
         data
@@ -128,7 +128,7 @@ def create_random_dataset_u8(nb, d, nq, k, path):
     from sklearn.datasets import make_blobs
     from sklearn.neighbors import NearestNeighbors
 
-    #print(f"Preparing datasets with {nb} random points and {nq} queries.")
+    # print(f"Preparing datasets with {nb} random points and {nq} queries.")
     os.mkdir(path)
     X, _ = make_blobs(n_samples=nb + nq, n_features=d, centers=nq, random_state=1)
 
@@ -145,7 +145,7 @@ def create_random_dataset_u8(nb, d, nq, k, path):
         np.array([nq, d], dtype="uint32").tofile(f)
         queries.tofile(f)
 
-    #print("Computing groundtruth")
+    # print("Computing groundtruth")
 
     nbrs = NearestNeighbors(n_neighbors=k, metric="euclidean", algorithm="brute").fit(
         data

apis/python/test/test_ingestion.py

@@ -8,9 +8,8 @@
 
 MINIMUM_ACCURACY = 0.9
 
-@pytest.mark.parametrize(
-  "query_type", ["heap", "nth"]
-)
+
+@pytest.mark.parametrize("query_type", ["heap", "nth"])
 def test_flat_ingestion_u8(tmp_path, query_type):
     dataset_dir = os.path.join(tmp_path, "dataset")
     array_uri = os.path.join(tmp_path, "array")
@@ -31,9 +30,8 @@ def test_flat_ingestion_u8(tmp_path, query_type):
     result = index.query(query_vectors, k=k, query_type=query_type)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-@pytest.mark.parametrize(
-  "query_type", ["heap", "nth"]
-)
+
+@pytest.mark.parametrize("query_type", ["heap", "nth"])
 def test_flat_ingestion_f32(tmp_path, query_type):
     dataset_dir = os.path.join(tmp_path, "dataset")
     array_uri = os.path.join(tmp_path, "array")
@@ -143,9 +141,7 @@ def test_ivf_flat_ingestion_f32(tmp_path):
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
-        query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL
-    )
+    result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
@@ -189,7 +185,5 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
-        query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL
-    )
+    result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY