Fix some TODOs after the cloud release (#505)

Fix some TODOs after the cloud release

Author: NikolaosPapailiou

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

@@ -1684,10 +1684,6 @@ def ingest_vectors_udf(
         trace_id: Optional[str] = None,
         distance_metric: vspy.DistanceMetric = vspy.DistanceMetric.SUM_OF_SQUARES,
     ):
-        import os
-        import random
-        import tempfile
-
         import tiledb.cloud
         from tiledb.vector_search.module import StdVector_u64
         from tiledb.vector_search.module import array_to_matrix
@@ -1702,11 +1698,6 @@ def ingest_vectors_udf(
         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_uri = partial_write_array_group[INDEX_ARRAY_NAME].uri
-        # Temporary solution until `ivf_index` library change gets released.
-        # TODO(nikos) remove this when the `partition_start` parameter of `ivf_index` gets released.
-        partial_write_array_index_array = tiledb.open(
-            partial_write_array_index_uri, "w", timestamp=index_timestamp
-        )
 
         for part in range(start, end, batch):
             part_end = part + batch
@@ -1715,32 +1706,7 @@ def ingest_vectors_udf(
 
             str(part) + "-" + str(part_end)
             part_id = int(part / batch)
-
-            # Temporary solution until `ivf_index` library change gets released. We create a local disk
-            # temporary array to hold the partial indices and write them to the respective range in the main array.
-            # TODO(nikos) remove this when the `partition_start` parameter of `ivf_index` gets released.
-            partial_write_array_index_tmp_uri = os.path.join(
-                tempfile.gettempdir(),
-                f"{random.randint(0,MAX_INT32)}_{part_id}",
-            )
-            index_array_rows_dim = tiledb.Dim(
-                name="rows",
-                domain=(0, MAX_INT32),
-                tile=100000,
-                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],
-            )
-            tiledb.Array.create(partial_write_array_index_tmp_uri, index_schema)
-            partition_start = part_id * (partitions + 1)
+            part_id * (partitions + 1)
 
             logger.debug("Input vectors start_pos: %d, end_pos: %d", part, part_end)
             updated_ids = read_updated_ids(
@@ -1758,12 +1724,11 @@ def ingest_vectors_udf(
                     deleted_ids=StdVector_u64(updated_ids),
                     centroids_uri=centroids_uri,
                     parts_uri=partial_write_array_parts_uri,
-                    index_array_uri=partial_write_array_index_tmp_uri,
-                    # index_array_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,
-                    # partition_start=part_id * (partitions + 1),
+                    partition_start=part_id * (partitions + 1),
                     nthreads=threads,
                     **(
                         {"timestamp": index_timestamp}
@@ -1801,12 +1766,11 @@ def ingest_vectors_udf(
                     deleted_ids=StdVector_u64(updated_ids),
                     centroids_uri=centroids_uri,
                     parts_uri=partial_write_array_parts_uri,
-                    index_array_uri=partial_write_array_index_tmp_uri,
-                    # index_array_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,
-                    # partition_start=part_id * (partitions + 1),
+                    partition_start=part_id * (partitions + 1),
                     nthreads=threads,
                     **(
                         {"timestamp": index_timestamp}
@@ -1816,15 +1780,6 @@ def ingest_vectors_udf(
                     config=config,
                 )
 
-            # Temporary solution until `ivf_index` library change gets released.
-            # TODO(nikos) remove this when the `partition_start` parameter of `ivf_index` gets released.
-            with tiledb.open(partial_write_array_index_tmp_uri) as a:
-                partial_write_array_index_array[
-                    partition_start : partition_start + partitions + 1
-                ] = a[0 : partitions + 1]
-            tiledb.Array.delete_array(partial_write_array_index_tmp_uri)
-        partial_write_array_index_array.close()
-
     def ingest_additions_udf(
         index_group_uri: str,
         updates_uri: str,
@@ -1837,10 +1792,6 @@ def ingest_additions_udf(
         verbose: bool = False,
         trace_id: Optional[str] = None,
     ):
-        import os
-        import random
-        import tempfile
-
         import tiledb.cloud
         from tiledb.vector_search.module import StdVector_u64
         from tiledb.vector_search.module import array_to_matrix
@@ -1855,30 +1806,6 @@ def ingest_additions_udf(
         partial_write_array_parts_uri = partial_write_array_group[PARTS_ARRAY_NAME].uri
         partial_write_array_index_uri = partial_write_array_group[INDEX_ARRAY_NAME].uri
 
-        # Temporary solution until `ivf_index` library change gets released. We create a local disk
-        # temporary array to hold the partial indices and write them to the respective range in the main array.
-        # TODO(nikos) remove this when the `partition_start` parameter of `ivf_index` gets released.
-        partial_write_array_index_tmp_uri = os.path.join(
-            tempfile.gettempdir(), f"{random.randint(0,MAX_INT32)}_{partition_start}"
-        )
-        index_array_rows_dim = tiledb.Dim(
-            name="rows",
-            domain=(0, MAX_INT32),
-            tile=100000,
-            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],
-        )
-        tiledb.Array.create(partial_write_array_index_tmp_uri, index_schema)
-
         additions_vectors, additions_external_ids = read_additions(
             updates_uri=updates_uri,
             config=config,
@@ -1904,29 +1831,16 @@ def ingest_additions_udf(
             deleted_ids=StdVector_u64(np.array([], np.uint64)),
             centroids_uri=centroids_uri,
             parts_uri=partial_write_array_parts_uri,
-            index_array_uri=partial_write_array_index_tmp_uri,
-            # index_array_uri=partial_write_array_index_uri,
+            index_array_uri=partial_write_array_index_uri,
             id_uri=partial_write_array_ids_uri,
             start=write_offset,
             end=0,
-            # partition_start=partition_start,
+            partition_start=partition_start,
             nthreads=threads,
             **({"timestamp": index_timestamp} if index_timestamp is not None else {}),
             config=config,
         )
 
-        # Temporary solution until `ivf_index` library change gets released.
-        # TODO(nikos) remove this when the `partition_start` parameter of `ivf_index` gets released.
-        partial_write_array_index_array = tiledb.open(
-            partial_write_array_index_uri, "w", timestamp=index_timestamp
-        )
-        with tiledb.open(partial_write_array_index_tmp_uri) as a:
-            partial_write_array_index_array[
-                partition_start : partition_start + partitions + 1
-            ] = a[0 : partitions + 1]
-        tiledb.Array.delete_array(partial_write_array_index_tmp_uri)
-        partial_write_array_index_array.close()
-
     def compute_partition_indexes_udf(
         index_group_uri: str,
         partitions: int,

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

@@ -427,8 +427,7 @@ def dist_qv_udf(
                 k_nn=k_nn,
                 ctx=Ctx(config),
                 timestamp=timestamp,
-                # TODO(nikos) add this after the library change gets released in TileDB cloud
-                # upper_bound=0 if memory_budget == -1 else memory_budget,
+                upper_bound=0 if memory_budget == -1 else memory_budget,
             )
             results = []
             for q in range(len(r)):