Format

Author: Nikos Papailiou

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

@@ -10,12 +10,14 @@
 IDS_ARRAY_NAME = "ids.tdb"
 PARTS_ARRAY_NAME = "parts.tdb"
 
+
 def submit_local(d, func, *args, **kwargs):
     # Drop kwarg
     kwargs.pop("image_name", None)
     kwargs.pop("resources", None)
     return d.submit_local(func, *args, **kwargs)
 
+
 class Index:
     def query(self, targets: np.ndarray, k=10, nqueries=10, nthreads=8, nprobe=1):
         raise NotImplementedError
@@ -205,7 +207,11 @@ def distributed_query(
         """
         from tiledb.cloud import dag
         from tiledb.cloud.dag import Mode
-        from tiledb.vector_search.module import array_to_matrix, partition_ivf_index, dist_qv
+        from tiledb.vector_search.module import (
+            array_to_matrix,
+            partition_ivf_index,
+            dist_qv,
+        )
         import math
         import numpy as np
         from functools import partial
@@ -218,7 +224,8 @@ def dist_qv_udf(
             active_partitions: np.array,
             active_queries: np.array,
             indices: np.array,
-            k_nn: int):
+            k_nn: int,
+        ):
             targets_m = array_to_matrix(query_vectors)
             r = dist_qv(
                 dtype=dtype,
@@ -264,10 +271,8 @@ def dist_qv_udf(
 
         targets_m = array_to_matrix(targets)
         active_partitions, active_queries = partition_ivf_index(
-            centroids=self._centroids,
-            query=targets_m,
-            nprobe=nprobe,
-            nthreads=nthreads)
+            centroids=self._centroids, query=targets_m, nprobe=nprobe, nthreads=nthreads
+        )
         num_parts = len(active_partitions)
 
         parts_per_node = int(math.ceil(num_parts / num_nodes))
@@ -276,19 +281,23 @@ def dist_qv_udf(
             part_end = part + parts_per_node
             if part_end > num_parts:
                 part_end = num_parts
-            nodes.append(submit(
-                dist_qv_udf,
-                dtype=self.dtype,
-                parts_uri=self.parts_db_uri,
-                ids_uri=self.ids_uri,
-                query_vectors=targets,
-                active_partitions=np.array(active_partitions)[part:part_end],
-                active_queries=np.array(active_queries[part:part_end], dtype=object),
-                indices=np.array(self._index),
-                k_nn=k,
-                resource_class='large',
-                image_name="3.9-vectorsearch",
-            ))
+            nodes.append(
+                submit(
+                    dist_qv_udf,
+                    dtype=self.dtype,
+                    parts_uri=self.parts_db_uri,
+                    ids_uri=self.ids_uri,
+                    query_vectors=targets,
+                    active_partitions=np.array(active_partitions)[part:part_end],
+                    active_queries=np.array(
+                        active_queries[part:part_end], dtype=object
+                    ),
+                    indices=np.array(self._index),
+                    k_nn=k,
+                    resource_class="large",
+                    image_name="3.9-vectorsearch",
+                )
+            )
 
         d.compute()
         d.wait()

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

@@ -422,7 +422,9 @@ def read_input_vectors(
         trace_id: Optional[str] = None,
     ) -> np.array:
         logger = setup(config, verbose)
-        logger.debug("Reading input vectors start_pos: %i, end_pos: %i", start_pos, end_pos)
+        logger.debug(
+            "Reading input vectors start_pos: %i, end_pos: %i", start_pos, end_pos
+        )
         if source_type == "TILEDB_ARRAY":
             with tiledb.open(source_uri, mode="r") as src_array:
                 return np.transpose(
@@ -511,7 +513,9 @@ def copy_centroids(
         logger = setup(config, verbose)
         group = tiledb.Group(array_uri)
         centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
-        logger.debug("Copying centroids from: %s, to: %s", copy_centroids_uri, centroids_uri)
+        logger.debug(
+            "Copying centroids from: %s, to: %s", copy_centroids_uri, centroids_uri
+        )
         src = tiledb.open(copy_centroids_uri, mode="r")
         dest = tiledb.open(centroids_uri, mode="w")
         src_centroids = src[:, :]
@@ -586,7 +590,9 @@ def init_centroids(
         trace_id: Optional[str] = None,
     ) -> np.array:
         logger = setup(config, verbose)
-        logger.debug("Initialising centroids by reading the first vectors in the source data.")
+        logger.debug(
+            "Initialising centroids by reading the first vectors in the source data."
+        )
         with tiledb.scope_ctx(ctx_or_config=config):
             return read_input_vectors(
                 source_uri=source_uri,
@@ -922,7 +928,9 @@ def consolidate_partition_udf(
     ):
         logger = setup(config, verbose)
         with tiledb.scope_ctx(ctx_or_config=config):
-            logger.debug("Consolidating partitions %d-%d", partition_id_start, partition_id_end)
+            logger.debug(
+                "Consolidating partitions %d-%d", partition_id_start, partition_id_end
+            )
             group = tiledb.Group(array_uri)
             partial_write_array_dir_uri = array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR
             partial_write_array_ids_uri = (
@@ -962,12 +970,16 @@ def consolidate_partition_udf(
             index_array = tiledb.open(index_array_uri, mode="r")
             ids_array = tiledb.open(ids_array_uri, mode="w")
             parts_array = tiledb.open(parts_array_uri, mode="w")
-            logger.debug("Partitions start: %d end: %d", partition_id_start, partition_id_end)
+            logger.debug(
+                "Partitions start: %d end: %d", partition_id_start, partition_id_end
+            )
             for part in range(partition_id_start, partition_id_end, batch):
                 part_end = part + batch
                 if part_end > partition_id_end:
                     part_end = partition_id_end
-                logger.debug("Consolidating partitions start: %d end: %d", part, part_end)
+                logger.debug(
+                    "Consolidating partitions start: %d end: %d", part, part_end
+                )
                 read_slices = []
                 for p in range(part, part_end):
                     for partition_slice in partition_slices[p]:
@@ -985,8 +997,13 @@ def consolidate_partition_udf(
                     "values"
                 ]
 
-                logger.debug("Ids shape %s, expected size: %d expected range:(%d,%d)", ids.shape, end_pos - start_pos,
-                             start_pos, end_pos)
+                logger.debug(
+                    "Ids shape %s, expected size: %d expected range:(%d,%d)",
+                    ids.shape,
+                    end_pos - start_pos,
+                    start_pos,
+                    end_pos,
+                )
                 if ids.shape[0] != end_pos - start_pos:
                     raise ValueError("Incorrect partition size.")
 
@@ -1338,7 +1355,10 @@ def consolidate_and_vacuum(
         logger.debug("input_vectors_per_work_item %d", input_vectors_per_work_item)
         logger.debug("input_vectors_work_items %d", input_vectors_work_items)
         logger.debug("input_vectors_work_tasks %d", input_vectors_work_tasks)
-        logger.debug("input_vectors_work_items_per_worker %d", input_vectors_work_items_per_worker)
+        logger.debug(
+            "input_vectors_work_items_per_worker %d",
+            input_vectors_work_items_per_worker,
+        )
 
         vectors_per_table_partitions = size / partitions
         table_partitions_per_work_item = int(
@@ -1354,10 +1374,15 @@ def consolidate_and_vacuum(
                 math.ceil(table_partitions_work_items / MAX_TASKS_PER_STAGE)
             )
             table_partitions_work_tasks = MAX_TASKS_PER_STAGE
-        logger.debug("table_partitions_per_work_item %d", table_partitions_per_work_item)
+        logger.debug(
+            "table_partitions_per_work_item %d", table_partitions_per_work_item
+        )
         logger.debug("table_partitions_work_items %d", table_partitions_work_items)
         logger.debug("table_partitions_work_tasks %d", table_partitions_work_tasks)
-        logger.debug("table_partitions_work_items_per_worker %d", table_partitions_work_items_per_worker)
+        logger.debug(
+            "table_partitions_work_items_per_worker %d",
+            table_partitions_work_items_per_worker,
+        )
 
         logger.debug("Creating arrays")
         create_arrays(

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

@@ -302,6 +302,7 @@ def partition_ivf_index(centroids, query, nprobe=1, nthreads=0):
     else:
         raise TypeError("Unsupported type!")
 
+
 def dist_qv(
     dtype: np.dtype,
     parts_uri: str,
@@ -311,7 +312,8 @@ def dist_qv(
     active_queries: np.array,
     indices: np.array,
     k_nn: int,
-    ctx: "Ctx" = None):
+    ctx: "Ctx" = None,
+):
     if ctx is None:
         ctx = Ctx({})
     args = tuple(
@@ -323,7 +325,7 @@ def dist_qv(
             active_queries,
             StdVector_u64(indices),
             ids_uri,
-            k_nn
+            k_nn,
         ]
     )
     if dtype == np.float32:
@@ -333,6 +335,7 @@ def dist_qv(
     else:
         raise TypeError("Unsupported type!")
 
+
 def validate_top_k(results: np.ndarray, ground_truth: np.ndarray):
     if results.dtype == np.uint64:
         return validate_top_k_u64(results, ground_truth)

apis/python/test/common.py

@@ -182,10 +182,11 @@ def create_array(path: str, data):
     with tiledb.open(path, "w") as A:
         A[:] = data
 
+
 def accuracy(result, gt):
     found = 0
     total = 0
     for i in range(len(result)):
-        total+=len(result[i])
-        found+=len(np.intersect1d(result[i], gt[i]))
-    return found/total
+        total += len(result[i])
+        found += len(np.intersect1d(result[i], gt[i]))
+    return found / total

apis/python/test/test_ingestion.py

@@ -6,6 +6,7 @@
 
 MINIMUM_ACCURACY = 0.9
 
+
 def test_flat_ingestion_u8(tmp_path):
     dataset_dir = os.path.join(tmp_path, "dataset")
     array_uri = os.path.join(tmp_path, "array")
@@ -47,6 +48,7 @@ def test_flat_ingestion_f32(tmp_path):
     result = np.transpose(index.query(np.transpose(query_vectors), k=k))
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
+
 def test_ivf_flat_ingestion_u8(tmp_path):
     dataset_dir = os.path.join(tmp_path, "dataset")
     array_uri = os.path.join(tmp_path, "array")
@@ -70,9 +72,7 @@ def test_ivf_flat_ingestion_u8(tmp_path):
         partitions=partitions,
         input_vectors_per_work_item=int(size / 10),
     )
-    result = np.transpose(
-        index.query(np.transpose(query_vectors), k=k, nprobe=nprobe)
-    )
+    result = np.transpose(index.query(np.transpose(query_vectors), k=k, nprobe=nprobe))
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
@@ -90,11 +90,18 @@ def test_ivf_flat_ingestion_u8(tmp_path):
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.distributed_query(np.transpose(query_vectors.astype(np.uint8)), k=k, nprobe=nprobe, mode=Mode.LOCAL)
+    result = index_ram.distributed_query(
+        np.transpose(query_vectors.astype(np.uint8)),
+        k=k,
+        nprobe=nprobe,
+        mode=Mode.LOCAL,
+    )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
+
 def test_ivf_flat_ingestion_f32(tmp_path):
     import time
+
     dataset_dir = os.path.join(tmp_path, "dataset")
     array_uri = os.path.join(tmp_path, "array")
     k = 10
@@ -120,9 +127,7 @@ def test_ivf_flat_ingestion_f32(tmp_path):
         input_vectors_per_work_item=int(size / 10),
     )
 
-    result = np.transpose(
-        index.query(np.transpose(query_vectors), k=k, nprobe=nprobe)
-    )
+    result = np.transpose(index.query(np.transpose(query_vectors), k=k, nprobe=nprobe))
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
@@ -140,11 +145,12 @@ def test_ivf_flat_ingestion_f32(tmp_path):
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.distributed_query(np.transpose(query_vectors), k=k, nprobe=nprobe, mode=Mode.LOCAL)
+    result = index_ram.distributed_query(
+        np.transpose(query_vectors), k=k, nprobe=nprobe, mode=Mode.LOCAL
+    )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
-
 def test_ivf_flat_ingestion_fvec(tmp_path):
     source_uri = "test/data/siftsmall/siftsmall_base.fvecs"
     queries_uri = "test/data/siftsmall/siftsmall_query.fvecs"
@@ -170,9 +176,7 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
         source_type=source_type,
         partitions=partitions,
     )
-    result = np.transpose(
-        index.query(np.transpose(query_vectors), k=k, nprobe=nprobe)
-    )
+    result = np.transpose(index.query(np.transpose(query_vectors), k=k, nprobe=nprobe))
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype)
@@ -190,5 +194,7 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.distributed_query(np.transpose(query_vectors), k=k, nprobe=nprobe, mode=Mode.LOCAL)
+    result = index_ram.distributed_query(
+        np.transpose(query_vectors), k=k, nprobe=nprobe, mode=Mode.LOCAL
+    )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY