Merge pull request #170 from TileDB-Inc/npapa/fix-update-uri-bug

Fix bug for creating the updates array using a tiledb URI

Author: NikolaosPapailiou

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

@@ -10,6 +10,7 @@
                                                   validate_storage_version)
 
 MAX_INT32 = np.iinfo(np.dtype("int32")).max
+MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 TILE_SIZE_BYTES = 128000000  # 128MB
 INDEX_TYPE = "FLAT"
 
@@ -139,8 +140,10 @@ def create(
         tile_size = TILE_SIZE_BYTES / np.dtype(vector_type).itemsize / dimensions
         ids_array_name = storage_formats[storage_version]["IDS_ARRAY_NAME"]
         parts_array_name = storage_formats[storage_version]["PARTS_ARRAY_NAME"]
+        updates_array_name = storage_formats[storage_version]["UPDATES_ARRAY_NAME"]
         ids_uri = f"{uri}/{ids_array_name}"
         parts_uri = f"{uri}/{parts_array_name}"
+        updates_array_uri = f"{uri}/{updates_array_name}"
 
         ids_array_rows_dim = tiledb.Dim(
             name="rows",
@@ -192,5 +195,21 @@ def create(
         tiledb.Array.create(parts_uri, parts_schema)
         group.add(parts_uri, name=parts_array_name)
 
+        external_id_dim = tiledb.Dim(
+            name="external_id",
+            domain=(0, MAX_UINT64 - 1),
+            dtype=np.dtype(np.uint64),
+        )
+        dom = tiledb.Domain(external_id_dim)
+        vector_attr = tiledb.Attr(name="vector", dtype=vector_type, var=True)
+        updates_schema = tiledb.ArraySchema(
+            domain=dom,
+            sparse=True,
+            attrs=[vector_attr],
+            allows_duplicates=False,
+        )
+        tiledb.Array.create(updates_array_uri, updates_schema)
+        group.add(updates_array_uri, name=updates_array_name)
+
         group.close()
         return FlatIndex(uri=uri, config=config)

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

@@ -52,9 +52,11 @@ def __init__(
             raise ValueError(
                 f"Time traveling is not supported for index storage_version={self.storage_version}"
             )
-
         updates_array_name = storage_formats[self.storage_version]["UPDATES_ARRAY_NAME"]
-        self.updates_array_uri = f"{self.group.uri}/{updates_array_name}"
+        if updates_array_name in self.group:
+            self.updates_array_uri = self.group[storage_formats[self.storage_version]["UPDATES_ARRAY_NAME"]].uri
+        else:
+            self.updates_array_uri = f"{self.group.uri}/{updates_array_name}"
         self.index_version = self.group.meta.get("index_version", "")
         self.ingestion_timestamps = [
             int(x)
@@ -134,7 +136,7 @@ def query(self, queries: np.ndarray, k, **kwargs):
             raise TypeError(f"A query in queries has {query_dimensions} dimensions, but the indexed data had {self.dimensions} dimensions")
 
         with tiledb.scope_ctx(ctx_or_config=self.config):
-            if not tiledb.array_exists(self.updates_array_uri):
+            if not tiledb.array_exists(self.updates_array_uri) or not self.has_updates():
                 if self.query_base_array:
                     return self.query_internal(queries, k, **kwargs)
                 else:
@@ -266,7 +268,22 @@ def get_dimensions(self):
     def query_internal(self, queries: np.ndarray, k, **kwargs):
         raise NotImplementedError
 
+    def has_updates(self):
+        if "has_updates" in self.group.meta:
+            return self.group.meta["has_updates"]
+        else:
+            return True
+
+    def set_has_updates(self, has_updates: bool = True):
+        if not self.group.meta["has_updates"]:
+            self.group.close()
+            self.group = tiledb.Group(self.uri, "w", ctx=tiledb.Ctx(self.config))
+            self.group.meta["has_updates"] = has_updates
+            self.group.close()
+            self.group = tiledb.Group(self.uri, "r", ctx=tiledb.Ctx(self.config))
+
     def update(self, vector: np.array, external_id: np.uint64, timestamp: int = None):
+        self.set_has_updates()
         updates_array = self.open_updates_array(timestamp=timestamp)
         vectors = np.empty((1), dtype="O")
         vectors[0] = vector
@@ -277,12 +294,14 @@ def update(self, vector: np.array, external_id: np.uint64, timestamp: int = None
     def update_batch(
         self, vectors: np.ndarray, external_ids: np.array, timestamp: int = None
     ):
+        self.set_has_updates()
         updates_array = self.open_updates_array(timestamp=timestamp)
         updates_array[external_ids] = {"vector": vectors}
         updates_array.close()
         self.consolidate_update_fragments()
 
     def delete(self, external_id: np.uint64, timestamp: int = None):
+        self.set_has_updates()
         updates_array = self.open_updates_array(timestamp=timestamp)
         deletes = np.empty((1), dtype="O")
         deletes[0] = np.array([], dtype=self.dtype)
@@ -291,6 +310,7 @@ def delete(self, external_id: np.uint64, timestamp: int = None):
         self.consolidate_update_fragments()
 
     def delete_batch(self, external_ids: np.array, timestamp: int = None):
+        self.set_has_updates()
         updates_array = self.open_updates_array(timestamp=timestamp)
         deletes = np.empty((len(external_ids)), dtype="O")
         for i in range(len(external_ids)):
@@ -529,4 +549,5 @@ def create_metadata(
         group.meta["index_type"] = index_type
         group.meta["base_sizes"] = json.dumps([0])
         group.meta["ingestion_timestamps"] = json.dumps([0])
+        group.meta["has_updates"] = False
         group.close()

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

@@ -12,6 +12,7 @@
                                                   validate_storage_version)
 
 MAX_INT32 = np.iinfo(np.dtype("int32")).max
+MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 TILE_SIZE_BYTES = 64000000  # 64MB
 INDEX_TYPE = "IVF_FLAT"
 
@@ -436,12 +437,8 @@ def dist_qv_udf(
             tmp = sorted(tmp_results, key=lambda t: t[0])[0:k]
             for j in range(len(tmp), k):
                 tmp.append((float(0.0), int(0)))
-            results_per_query_d.append(
-                np.array(tmp, dtype=np.dtype("float,uint64"))["f0"]
-            )
-            results_per_query_i.append(
-                np.array(tmp, dtype=np.dtype("float,uint64"))["f1"]
-            )
+            results_per_query_d.append(np.array(tmp, dtype=np.float32)[:, 0])
+            results_per_query_i.append(np.array(tmp, dtype=np.uint64)[:, 1])
         return np.array(results_per_query_d), np.array(results_per_query_i)
 
 
@@ -474,10 +471,12 @@ def create(
         index_array_name = storage_formats[storage_version]["INDEX_ARRAY_NAME"]
         ids_array_name = storage_formats[storage_version]["IDS_ARRAY_NAME"]
         parts_array_name = storage_formats[storage_version]["PARTS_ARRAY_NAME"]
+        updates_array_name = storage_formats[storage_version]["UPDATES_ARRAY_NAME"]
         centroids_uri = f"{uri}/{centroids_array_name}"
         index_array_uri = f"{uri}/{index_array_name}"
         ids_uri = f"{uri}/{ids_array_name}"
         parts_uri = f"{uri}/{parts_array_name}"
+        updates_array_uri = f"{uri}/{updates_array_name}"
 
         centroids_array_rows_dim = tiledb.Dim(
             name="rows",
@@ -581,5 +580,21 @@ def create(
         tiledb.Array.create(parts_uri, parts_schema)
         group.add(parts_uri, name=parts_array_name)
 
+        external_id_dim = tiledb.Dim(
+            name="external_id",
+            domain=(0, MAX_UINT64 - 1),
+            dtype=np.dtype(np.uint64),
+        )
+        dom = tiledb.Domain(external_id_dim)
+        vector_attr = tiledb.Attr(name="vector", dtype=vector_type, var=True)
+        updates_schema = tiledb.ArraySchema(
+            domain=dom,
+            sparse=True,
+            attrs=[vector_attr],
+            allows_duplicates=False,
+        )
+        tiledb.Array.create(updates_array_uri, updates_schema)
+        group.add(updates_array_uri, name=updates_array_name)
+
         group.close()
         return IVFFlatIndex(uri=uri, config=config, memory_budget=1000000)

apis/python/test/test_cloud.py

@@ -50,9 +50,17 @@ def test_cloud_flat(self):
             config=tiledb.cloud.Config().dict(),
             mode=Mode.BATCH,
         )
+        tiledb_index_uri = groups.info(index_uri).tiledb_uri
+        index = vs.flat_index.FlatIndex(uri=tiledb_index_uri)
+
         _, result_i = index.query(queries, k=k)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
+        index.delete(external_id=42)
+        _, result_i = index.query(queries, k=k)
+        assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
+
+
     def test_cloud_ivf_flat(self):
         source_uri = "tiledb://TileDB-Inc/sift_10k"
         queries_uri = "test/data/siftsmall/siftsmall_query.fvecs"
@@ -79,6 +87,9 @@ def test_cloud_ivf_flat(self):
             # mode=Mode.BATCH,
         )
 
+        tiledb_index_uri = groups.info(index_uri).tiledb_uri
+        index = vs.ivf_flat_index.IVFFlatIndex(uri=tiledb_index_uri)
+
         _, result_i = index.query(queries, k=k, nprobe=nprobe)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
@@ -113,4 +124,8 @@ def test_cloud_ivf_flat(self):
         with self.assertRaises(TypeError):
             index.query(queries, k=k, nprobe=nprobe, mode=Mode.REALTIME, resource_class="large", resources=resources)
         with self.assertRaises(TypeError):
-            index.query(queries, k=k, nprobe=nprobe, mode=Mode.BATCH, resource_class="large", resources=resources)
+            index.query(queries, k=k, nprobe=nprobe, mode=Mode.BATCH, resource_class="large", resources=resources)
+
+        index.delete(external_id=42)
+        _, result_i = index.query(queries, k=k, nprobe=nprobe)
+        assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY