Make vector ingestion work with tiledb:// uris

Author: Nikos Papailiou

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

@@ -39,12 +39,15 @@ class FlatIndex(Index):
         Optional name of partitions
     """
 
-    def __init__(self, uri: str, dtype: np.dtype, parts_name: str = "parts.tdb"):
+    def __init__(self, uri: str, dtype: np.dtype, parts_name: str = "parts.tdb", ctx: "Ctx" = None):
         self.uri = uri
         self.dtype = dtype
         self._index = None
+        self.ctx = ctx
+        if ctx is None:
+            self.ctx = Ctx({})
 
-        self._db = load_as_matrix(os.path.join(uri, parts_name))
+        self._db = load_as_matrix(os.path.join(uri, parts_name), ctx=self.ctx)
 
     def query(
         self,
@@ -118,10 +121,10 @@ def __init__(
 
         # TODO pass in a context
         if self.memory_budget == -1:
-            self._db = load_as_matrix(self.parts_db_uri)
+            self._db = load_as_matrix(self.parts_db_uri, ctx=self.ctx)
             self._ids = read_vector_u64(self.ctx, self.ids_uri)
 
-        self._centroids = load_as_matrix(self.centroids_uri)
+        self._centroids = load_as_matrix(self.centroids_uri, ctx=self.ctx)
         self._index = read_vector_u64(self.ctx, self.index_uri)
 
     def query(

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

@@ -223,7 +223,7 @@ def create_arrays(
                 )
                 logger.debug(parts_schema)
                 tiledb.Array.create(parts_uri, parts_schema)
-                group.add(PARTS_ARRAY_NAME, name=PARTS_ARRAY_NAME, relative=True)
+                group.add(parts_uri, name=PARTS_ARRAY_NAME)
 
         elif index_type == "IVF_FLAT":
             centroids_uri = f"{group.uri}/{CENTROIDS_ARRAY_NAME}"
@@ -272,7 +272,7 @@ def create_arrays(
                 logger.debug(centroids_schema)
                 tiledb.Array.create(centroids_uri, centroids_schema)
                 group.add(
-                    CENTROIDS_ARRAY_NAME, name=CENTROIDS_ARRAY_NAME, relative=True
+                    centroids_uri, name=CENTROIDS_ARRAY_NAME
                 )
 
             if not tiledb.array_exists(index_uri):
@@ -295,7 +295,7 @@ def create_arrays(
                 )
                 logger.debug(index_schema)
                 tiledb.Array.create(index_uri, index_schema)
-                group.add(INDEX_ARRAY_NAME, name=INDEX_ARRAY_NAME, relative=True)
+                group.add(index_uri, name=INDEX_ARRAY_NAME)
 
             if not tiledb.array_exists(ids_uri):
                 logger.debug("Creating ids array")
@@ -317,7 +317,7 @@ def create_arrays(
                 )
                 logger.debug(ids_schema)
                 tiledb.Array.create(ids_uri, ids_schema)
-                group.add(IDS_ARRAY_NAME, name=IDS_ARRAY_NAME, relative=True)
+                group.add(ids_uri, name=IDS_ARRAY_NAME)
 
             if not tiledb.array_exists(parts_uri):
                 logger.debug("Creating parts array")
@@ -347,15 +347,27 @@ def create_arrays(
                 )
                 logger.debug(parts_schema)
                 tiledb.Array.create(parts_uri, parts_schema)
-                group.add(PARTS_ARRAY_NAME, name=PARTS_ARRAY_NAME, relative=True)
-
-            vfs = tiledb.VFS()
-            if vfs.is_dir(partial_write_array_dir_uri):
-                vfs.remove_dir(partial_write_array_dir_uri)
-            vfs.create_dir(partial_write_array_dir_uri)
-            if vfs.is_dir(partial_write_array_index_uri):
-                vfs.remove_dir(partial_write_array_index_uri)
-            vfs.create_dir(partial_write_array_index_uri)
+                group.add(parts_uri, name=PARTS_ARRAY_NAME)
+
+            try:
+                tiledb.group_create(partial_write_array_dir_uri)
+            except tiledb.TileDBError as err:
+                message = str(err)
+                if "already exists" in message:
+                    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)
+
+            try:
+                tiledb.group_create(partial_write_array_index_uri)
+            except tiledb.TileDBError as err:
+                message = str(err)
+                if "already exists" in message:
+                    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")
 
             if not tiledb.array_exists(partial_write_array_ids_uri):
                 logger.debug("Creating temp ids array")
@@ -377,6 +389,7 @@ 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)
 
             if not tiledb.array_exists(partial_write_array_parts_uri):
                 logger.debug("Creating temp parts array")
@@ -407,6 +420,34 @@ 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)
+
+            for part in range(1):
+                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(
+                        name="rows",
+                        domain=(0, partitions),
+                        tile=partitions,
+                        dtype=np.dtype(np.int32),
+                    )
+                    index_array_dom = tiledb.Domain(index_array_rows_dim)
+                    index_attr = tiledb.Attr(name="values", dtype=np.dtype(np.uint64))
+                    index_schema = tiledb.ArraySchema(
+                        domain=index_array_dom,
+                        sparse=False,
+                        attrs=[index_attr],
+                        capacity=partitions,
+                        cell_order="col-major",
+                        tile_order="col-major",
+                    )
+                    logger.debug(index_schema)
+                    tiledb.Array.create(part_index_uri, index_schema)
+                    partial_write_array_index_group.add(part_index_uri, name=str(part))
+            partial_write_array_group.close()
+            partial_write_array_index_group.close()
+
         else:
             raise ValueError(f"Not supported index_type {index_type}")
 
@@ -822,21 +863,21 @@ def ingest_vectors_udf(
         logger = setup(config, verbose)
         group = tiledb.Group(array_uri)
         centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
-        partial_write_array_dir_uri = array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR
-        partial_write_array_ids_uri = partial_write_array_dir_uri + "/" + IDS_ARRAY_NAME
-        partial_write_array_parts_uri = (
-            partial_write_array_dir_uri + "/" + PARTS_ARRAY_NAME
-        )
+        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)
 
         for part in range(start, end, batch):
             part_end = part + batch
             if part_end > end:
                 part_end = end
 
             part_name = str(part) + "-" + str(part_end)
-            partial_write_array_index_uri = (
-                partial_write_array_dir_uri + "/" + INDEX_ARRAY_NAME + "/" + part_name
-            )
+
+            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")
@@ -889,12 +930,14 @@ def compute_partition_indexes_udf(
         with tiledb.scope_ctx(ctx_or_config=config):
             group = tiledb.Group(array_uri)
             index_array_uri = group[INDEX_ARRAY_NAME].uri
-            vfs = tiledb.VFS()
+            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)
             partition_sizes = np.zeros(partitions)
             indexes = np.zeros(partitions + 1).astype(np.uint64)
-            for partial_index_array_uri in vfs.ls(
-                array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR + "/" + INDEX_ARRAY_NAME
-            ):
+            for part in partial_write_array_index_group:
+                partial_index_array_uri = part.uri
                 if tiledb.array_exists(partial_index_array_uri):
                     partial_index_array = tiledb.open(partial_index_array_uri, mode="r")
                     partial_indexes = partial_index_array[:]["values"]
@@ -912,7 +955,7 @@ def compute_partition_indexes_udf(
                 _sum += partition_size
                 i += 1
             indexes[i] = _sum
-            logger.debug("Partition indexes: %d", indexes)
+            logger.debug(f"Partition indexes: {indexes}")
             index_array = tiledb.open(index_array_uri, mode="w")
             index_array[:] = indexes
 
@@ -932,23 +975,21 @@ def consolidate_partition_udf(
                 "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 = (
-                partial_write_array_dir_uri + "/" + IDS_ARRAY_NAME
-            )
-            partial_write_array_parts_uri = (
-                partial_write_array_dir_uri + "/" + PARTS_ARRAY_NAME
-            )
+            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)
             index_array_uri = group[INDEX_ARRAY_NAME].uri
             ids_array_uri = group[IDS_ARRAY_NAME].uri
             parts_array_uri = group[PARTS_ARRAY_NAME].uri
             vfs = tiledb.VFS()
             partition_slices = []
             for i in range(partitions):
                 partition_slices.append([])
-            for partial_index_array_uri in vfs.ls(
-                array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR + "/" + INDEX_ARRAY_NAME
-            ):
+            for part in partial_write_array_index_group:
+                partial_index_array_uri = part.uri
                 if tiledb.array_exists(partial_index_array_uri):
                     partial_index_array = tiledb.open(partial_index_array_uri, mode="r")
                     partial_indexes = partial_index_array[:]["values"]
@@ -1294,12 +1335,14 @@ def consolidate_and_vacuum(
                 tiledb.consolidate(group[IDS_ARRAY_NAME].uri, config=conf)
                 tiledb.vacuum(group[IDS_ARRAY_NAME].uri, config=conf)
 
-        vfs = tiledb.VFS(config)
-        partial_write_array_dir_uri = array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR
-        if vfs.is_dir(partial_write_array_dir_uri):
-            vfs.remove_dir(partial_write_array_dir_uri)
+        # TODO remove temp data for tiledb URIs
+        if not array_uri.startswith("tiledb://"):
+            vfs = tiledb.VFS(config)
+            partial_write_array_dir_uri = array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR
+            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):
+    with tiledb.scope_ctx(ctx_or_config=config) as ctx:
         logger = setup(config, verbose)
         logger.debug("Ingesting Vectors into %r", array_uri)
         try:
@@ -1419,6 +1462,6 @@ 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)
+            return FlatIndex(uri=array_uri, dtype=vector_type, ctx=ctx)
         elif index_type == "IVF_FLAT":
-            return IVFFlatIndex(uri=array_uri, dtype=vector_type, memory_budget=1000000)
+            return IVFFlatIndex(uri=array_uri, dtype=vector_type, memory_budget=1000000, ctx=ctx)

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

@@ -8,7 +8,7 @@
 from typing import Optional
 
 
-def load_as_matrix(path: str, nqueries: int = 0, config: Dict = {}):
+def load_as_matrix(path: str, nqueries: int = 0, ctx: "Ctx" = None):
     """
     Load array as Matrix class
 
@@ -21,7 +21,8 @@ def load_as_matrix(path: str, nqueries: int = 0, config: Dict = {}):
     config: Dict
         TileDB configuration parameters
     """
-    ctx = Ctx(config)
+    if ctx is None:
+        ctx = Ctx({})
 
     a = tiledb.ArraySchema.load(path)
     dtype = a.attr(0).dtype
@@ -43,7 +44,7 @@ def load_as_matrix(path: str, nqueries: int = 0, config: Dict = {}):
     return m
 
 
-def load_as_array(path, return_matrix: bool = False, config: Dict = {}):
+def load_as_array(path, return_matrix: bool = False,  ctx: "Ctx" = None):
     """
     Load array as array class
 
@@ -56,7 +57,7 @@ def load_as_array(path, return_matrix: bool = False, config: Dict = {}):
     config: Dict
         TileDB configuration parameters
     """
-    m = load_as_matrix(path, config=config)
+    m = load_as_matrix(path, ctx=ctx)
     r = np.array(m, copy=False)
 
     # hang on to a copy for testing purposes, for now

src/include/detail/ivf/index.h

@@ -136,7 +136,7 @@ int ivf_index(
           ctx, shuffled_db, parts_uri, start_pos, false);
     }
     if (index_uri != "") {
-      write_vector<ids_type>(ctx, indices, index_uri, 0, true);
+      write_vector<ids_type>(ctx, indices, index_uri, 0, false);
     }
     if (id_uri != "") {
       write_vector<ids_type>(ctx, shuffled_ids, id_uri, start_pos, false);