Rename array_uri -> index_uri, and clarify internal variable names

Author: ihnorton

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

@@ -134,7 +134,7 @@ def __init__(
         self.centroids_uri = group[
             storage_formats[self.storage_version]["CENTROIDS_ARRAY_NAME"]
         ].uri
-        self.index_uri = group[
+        self.index_array_uri = group[
             storage_formats[self.storage_version]["INDEX_ARRAY_NAME"]
         ].uri
         self.ids_uri = group[
@@ -145,7 +145,7 @@ def __init__(
         self._centroids = load_as_matrix(
             self.centroids_uri, ctx=self.ctx, config=config
         )
-        self._index = read_vector_u64(self.ctx, self.index_uri)
+        self._index = read_vector_u64(self.ctx, self.index_array_uri)
 
         # TODO pass in a context
         if self.memory_budget == -1:

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

@@ -7,7 +7,7 @@
 
 def ingest(
     index_type: str,
-    array_uri: str,
+    index_uri: str,
     source_uri: str,
     source_type: str,
     *,
@@ -30,8 +30,8 @@ def ingest(
     ----------
     index_type: str
         Type of vector index (FLAT, IVF_FLAT)
-    array_uri: str
-        Vector array URI
+    index_uri: str
+        Vector index URI (stored as TileDB group)
     source_uri: str
         Data source URI
     source_type: str
@@ -81,6 +81,9 @@ def ingest(
     from tiledb.cloud.utilities import set_aws_context
     from tiledb.vector_search.storage_formats import storage_formats, STORAGE_VERSION
 
+    # use index_group_uri for internal clarity
+    index_group_uri = index_uri
+
     CENTROIDS_ARRAY_NAME = storage_formats[STORAGE_VERSION]["CENTROIDS_ARRAY_NAME"]
     INDEX_ARRAY_NAME = storage_formats[STORAGE_VERSION]["INDEX_ARRAY_NAME"]
     IDS_ARRAY_NAME = storage_formats[STORAGE_VERSION]["IDS_ARRAY_NAME"]
@@ -232,7 +235,7 @@ def create_arrays(
 
         elif index_type == "IVF_FLAT":
             centroids_uri = f"{group.uri}/{CENTROIDS_ARRAY_NAME}"
-            index_uri = f"{group.uri}/{INDEX_ARRAY_NAME}"
+            index_array_uri = f"{group.uri}/{INDEX_ARRAY_NAME}"
             ids_uri = f"{group.uri}/{IDS_ARRAY_NAME}"
             parts_uri = f"{group.uri}/{PARTS_ARRAY_NAME}"
             partial_write_array_dir_uri = f"{group.uri}/{PARTIAL_WRITE_ARRAY_DIR}"
@@ -280,7 +283,7 @@ def create_arrays(
                 tiledb.Array.create(centroids_uri, centroids_schema)
                 group.add(centroids_uri, name=CENTROIDS_ARRAY_NAME)
 
-            if not tiledb.array_exists(index_uri):
+            if not tiledb.array_exists(index_array_uri):
                 logger.debug("Creating index array")
                 index_array_rows_dim = tiledb.Dim(
                     name="rows",
@@ -303,8 +306,8 @@ def create_arrays(
                     tile_order="col-major",
                 )
                 logger.debug(index_schema)
-                tiledb.Array.create(index_uri, index_schema)
-                group.add(index_uri, name=INDEX_ARRAY_NAME)
+                tiledb.Array.create(index_array_uri, index_schema)
+                group.add(index_array_uri, name=INDEX_ARRAY_NAME)
 
             if not tiledb.array_exists(ids_uri):
                 logger.debug("Creating ids array")
@@ -582,14 +585,14 @@ def read_input_vectors(
     # --------------------------------------------------------------------
 
     def copy_centroids(
-        array_uri: str,
+        index_group_uri: str,
         copy_centroids_uri: str,
         config: Optional[Mapping[str, Any]] = None,
         verbose: bool = False,
         trace_id: Optional[str] = None,
     ):
         logger = setup(config, verbose)
-        group = tiledb.Group(array_uri)
+        group = tiledb.Group(index_group_uri)
         centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
         logger.debug(
             "Copying centroids from: %s, to: %s", copy_centroids_uri, centroids_uri
@@ -604,7 +607,7 @@ def copy_centroids(
     # centralised kmeans UDFs
     # --------------------------------------------------------------------
     def centralised_kmeans(
-        array_uri: str,
+        index_group_uri: str,
         source_uri: str,
         source_type: str,
         vector_type: np.dtype,
@@ -623,7 +626,7 @@ def centralised_kmeans(
 
         with tiledb.scope_ctx(ctx_or_config=config):
             logger = setup(config, verbose)
-            group = tiledb.Group(array_uri)
+            group = tiledb.Group(index_group_uri)
             centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
             sample_vectors = read_input_vectors(
                 source_uri=source_uri,
@@ -811,7 +814,7 @@ def compute_new_centroids(*argv):
         return np.mean(argv, axis=0).astype(np.float32)
 
     def ingest_flat(
-        array_uri: str,
+        index_group_uri: str,
         source_uri: str,
         source_type: str,
         vector_type: np.dtype,
@@ -829,7 +832,7 @@ def ingest_flat(
 
         logger = setup(config, verbose)
         with tiledb.scope_ctx(ctx_or_config=config):
-            group = tiledb.Group(array_uri)
+            group = tiledb.Group(index_group_uri)
             parts_array_uri = group[PARTS_ARRAY_NAME].uri
             target = tiledb.open(parts_array_uri, mode="w")
             logger.debug("Input vectors start_pos: %d, end_pos: %d", start, end)
@@ -857,7 +860,7 @@ def ingest_flat(
 
     def write_centroids(
         centroids: np.array,
-        array_uri: str,
+        index_group_uri: str,
         partitions: int,
         dimensions: int,
         config: Optional[Mapping[str, Any]] = None,
@@ -866,7 +869,7 @@ def write_centroids(
     ):
         with tiledb.scope_ctx(ctx_or_config=config):
             logger = setup(config, verbose)
-            group = tiledb.Group(array_uri)
+            group = tiledb.Group(index_group_uri)
             centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
             logger.debug("Writing centroids to array %s", centroids_uri)
             with tiledb.open(centroids_uri, mode="w") as A:
@@ -876,7 +879,7 @@ def write_centroids(
     # vector ingestion UDFs
     # --------------------------------------------------------------------
     def ingest_vectors_udf(
-        array_uri: str,
+        index_group_uri: str,
         source_uri: str,
         source_type: str,
         vector_type: np.dtype,
@@ -898,7 +901,7 @@ def ingest_vectors_udf(
         )
 
         logger = setup(config, verbose)
-        group = tiledb.Group(array_uri)
+        group = tiledb.Group(index_uri)
         centroids_uri = group[CENTROIDS_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)
@@ -929,7 +932,7 @@ def ingest_vectors_udf(
                     db_uri=source_uri,
                     centroids_uri=centroids_uri,
                     parts_uri=partial_write_array_parts_uri,
-                    index_uri=partial_write_array_index_uri,
+                    index_array_uri=partial_write_array_index_uri,
                     id_uri=partial_write_array_ids_uri,
                     start=part,
                     end=part_end,
@@ -954,7 +957,7 @@ def ingest_vectors_udf(
                     db=array_to_matrix(np.transpose(in_vectors).astype(vector_type)),
                     centroids_uri=centroids_uri,
                     parts_uri=partial_write_array_parts_uri,
-                    index_uri=partial_write_array_index_uri,
+                    index_array_uri=partial_write_array_index_uri,
                     id_uri=partial_write_array_ids_uri,
                     start=part,
                     end=part_end,
@@ -963,15 +966,15 @@ def ingest_vectors_udf(
                 )
 
     def compute_partition_indexes_udf(
-        array_uri: str,
+        index_group_uri: str,
         partitions: int,
         config: Optional[Mapping[str, Any]] = None,
         verbose: bool = False,
         trace_id: Optional[str] = None,
     ):
         logger = setup(config, verbose)
         with tiledb.scope_ctx(ctx_or_config=config):
-            group = tiledb.Group(array_uri)
+            group = tiledb.Group(index_group_uri)
             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)
@@ -1007,7 +1010,7 @@ def compute_partition_indexes_udf(
             index_array[:] = indexes
 
     def consolidate_partition_udf(
-        array_uri: str,
+        index_group_uri: str,
         partition_id_start: int,
         partition_id_end: int,
         batch: int,
@@ -1021,7 +1024,7 @@ def consolidate_partition_udf(
             logger.debug(
                 "Consolidating partitions %d-%d", partition_id_start, partition_id_end
             )
-            group = tiledb.Group(array_uri)
+            group = tiledb.Group(index_group_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_ids_uri = partial_write_array_group[IDS_ARRAY_NAME].uri
@@ -1119,7 +1122,7 @@ def submit_local(d, func, *args, **kwargs):
 
     def create_ingestion_dag(
         index_type: str,
-        array_uri: str,
+        index_group_uri: str,
         source_uri: str,
         source_type: str,
         vector_type: np.dtype,
@@ -1174,7 +1177,7 @@ def create_ingestion_dag(
                     end = size
                 ingest_node = submit(
                     ingest_flat,
-                    array_uri=array_uri,
+                    index_group_uri=index_group_uri,
                     source_uri=source_uri,
                     source_type=source_type,
                     vector_type=vector_type,
@@ -1195,7 +1198,7 @@ def create_ingestion_dag(
             if copy_centroids_uri is not None:
                 centroids_node = submit(
                     copy_centroids,
-                    array_uri=array_uri,
+                    index_group_uri=index_group_uri,
                     copy_centroids_uri=copy_centroids_uri,
                     config=config,
                     verbose=verbose,
@@ -1208,7 +1211,7 @@ def create_ingestion_dag(
                 if training_sample_size <= CENTRALISED_KMEANS_MAX_SAMPLE_SIZE:
                     centroids_node = submit(
                         centralised_kmeans,
-                        array_uri=array_uri,
+                        index_group_uri=index_group_uri,
                         source_uri=source_uri,
                         source_type=source_type,
                         vector_type=vector_type,
@@ -1291,7 +1294,7 @@ def create_ingestion_dag(
                     centroids_node = submit(
                         write_centroids,
                         centroids=internal_centroids_node,
-                        array_uri=array_uri,
+                        index_group_uri=index_group_uri,
                         partitions=partitions,
                         dimensions=dimensions,
                         config=config,
@@ -1304,7 +1307,7 @@ def create_ingestion_dag(
 
             compute_indexes_node = submit(
                 compute_partition_indexes_udf,
-                array_uri=array_uri,
+                index_group_uri=index_group_uri,
                 partitions=partitions,
                 config=config,
                 verbose=verbose,
@@ -1322,7 +1325,7 @@ def create_ingestion_dag(
                     end = size
                 ingest_node = submit(
                     ingest_vectors_udf,
-                    array_uri=array_uri,
+                    index_group_uri=index_group_uri,
                     source_uri=source_uri,
                     source_type=source_type,
                     vector_type=vector_type,
@@ -1354,7 +1357,7 @@ def create_ingestion_dag(
                     end = partitions
                 consolidate_partition_node = submit(
                     consolidate_partition_udf,
-                    array_uri=array_uri,
+                    index_group_uri=index_group_uri,
                     partition_id_start=start,
                     partition_id_end=end,
                     batch=table_partitions_per_work_item,
@@ -1373,39 +1376,45 @@ def create_ingestion_dag(
             raise ValueError(f"Not supported index_type {index_type}")
 
     def consolidate_and_vacuum(
-        array_uri: str,
+        index_group_uri: str,
         config: Optional[Mapping[str, Any]] = None,
     ):
         modes = ["fragment_meta", "commits", "array_meta"]
         for mode in modes:
             conf = tiledb.Config(config)
             conf["sm.consolidation.mode"] = mode
             conf["sm.vacuum.mode"] = mode
-            group = tiledb.Group(array_uri, config=conf)
+            group = tiledb.Group(index_group_uri, config=conf)
             tiledb.consolidate(group[PARTS_ARRAY_NAME].uri, config=conf)
             tiledb.vacuum(group[PARTS_ARRAY_NAME].uri, config=conf)
             if index_type == "IVF_FLAT":
                 tiledb.consolidate(group[IDS_ARRAY_NAME].uri, config=conf)
                 tiledb.vacuum(group[IDS_ARRAY_NAME].uri, config=conf)
 
         # TODO remove temp data for tiledb URIs
-        if not array_uri.startswith("tiledb://"):
+        if not index_group_uri.startswith("tiledb://"):
             vfs = tiledb.VFS(config)
-            partial_write_array_dir_uri = array_uri + "/" + PARTIAL_WRITE_ARRAY_DIR
+            partial_write_array_dir_uri = index_group_uri + "/" + PARTIAL_WRITE_ARRAY_DIR
             if vfs.is_dir(partial_write_array_dir_uri):
                 vfs.remove_dir(partial_write_array_dir_uri)
 
+
+    # --------------------------------------------------------------------
+    # End internal function definitions
+    # --------------------------------------------------------------------
+
+
     with tiledb.scope_ctx(ctx_or_config=config):
         logger = setup(config, verbose)
-        logger.debug("Ingesting Vectors into %r", array_uri)
+        logger.debug("Ingesting Vectors into %r", index_uri)
         try:
-            tiledb.group_create(array_uri)
+            tiledb.group_create(index_group_uri)
         except tiledb.TileDBError as err:
             message = str(err)
             if "already exists" in message:
-                logger.debug(f"Group '{array_uri}' already exists")
+                logger.debug(f"Group '{index_group_uri}' already exists")
             raise err
-        group = tiledb.Group(array_uri, "w")
+        group = tiledb.Group(index_group_uri, "w")
 
         in_size, dimensions, vector_type = read_source_metadata(
             source_uri=source_uri, source_type=source_type, logger=logger
@@ -1492,7 +1501,7 @@ def consolidate_and_vacuum(
         logger.debug("Creating ingestion graph")
         d = create_ingestion_dag(
             index_type=index_type,
-            array_uri=array_uri,
+            index_group_uri=index_group_uri,
             source_uri=source_uri,
             source_type=source_type,
             vector_type=vector_type,
@@ -1515,9 +1524,9 @@ def consolidate_and_vacuum(
         d.compute()
         logger.debug("Submitted ingestion graph")
         d.wait()
-        consolidate_and_vacuum(array_uri=array_uri, config=config)
+        consolidate_and_vacuum(index_group_uri=index_group_uri, config=config)
 
         if index_type == "FLAT":
-            return FlatIndex(uri=array_uri, config=config)
+            return FlatIndex(uri=index_group_uri, config=config)
         elif index_type == "IVF_FLAT":
-            return IVFFlatIndex(uri=array_uri, memory_budget=1000000, config=config)
+            return IVFFlatIndex(uri=index_group_uri, memory_budget=1000000, config=config)