Clean-up ingestion.py and reduce non-verbose log messages. (#82)

* Fix double imports.

* Log the schemas of the created arrays in debug mode.

* Use percent formatting for the log messages.

* Do not redefine the built-in function `sum`.

* Downgrade many log messages to debug severity.

Author: teo-tsirpanis

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

@@ -69,15 +69,14 @@ def ingest(
     import enum
     import logging
     import math
-    from typing import Any, Mapping, Optional
+    from typing import Any, Mapping
     import multiprocessing
     import os
 
     import numpy as np
 
     import tiledb
     from tiledb.cloud import dag
-    from tiledb.cloud.dag import Mode
     from tiledb.cloud.rest_api import models
     from tiledb.cloud.utilities import get_logger
     from tiledb.cloud.utilities import set_aws_context
@@ -423,7 +422,7 @@ def read_input_vectors(
         trace_id: Optional[str] = None,
     ) -> np.array:
         logger = setup(config, verbose)
-        logger.info(f"Reading input vectors start_pos: {start_pos}, end_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(
@@ -512,9 +511,7 @@ def copy_centroids(
         logger = setup(config, verbose)
         group = tiledb.Group(array_uri)
         centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
-        logger.info(
-            f"Copying centroids from: {copy_centroids_uri}, to: {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[:, :]
@@ -560,7 +557,7 @@ def centralised_kmeans(
             verb = 0
             if verbose:
                 verb = 3
-            logger.info("Start kmeans training")
+            logger.debug("Start kmeans training")
             km = KMeans(
                 n_clusters=partitions,
                 init=init,
@@ -569,7 +566,7 @@ def centralised_kmeans(
                 n_init=n_init,
             )
             km.fit_predict(sample_vectors)
-            logger.info(f"Writing centroids to array {centroids_uri}")
+            logger.debug("Writing centroids to array %s", centroids_uri)
             with tiledb.open(centroids_uri, mode="w") as A:
                 A[0:dimensions, 0:partitions] = np.transpose(
                     np.array(km.cluster_centers_)
@@ -589,9 +586,7 @@ def init_centroids(
         trace_id: Optional[str] = None,
     ) -> np.array:
         logger = setup(config, verbose)
-        logger.info(
-            "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,
@@ -632,7 +627,7 @@ def generate_new_centroid_per_thread(
             new_centroid_count = np.ones(len(cents_t))
             for vector_id in range(start, end):
                 if vector_id % 100000 == 0:
-                    logger.debug(f"Vectors computed: {vector_id}")
+                    logger.debug("Vectors computed: %d", vector_id)
                 c_id = assignments_t[vector_id]
                 if new_centroid_count[c_id] == 1:
                     new_centroid_sums[c_id] = vectors_t[vector_id]
@@ -642,13 +637,13 @@ def generate_new_centroid_per_thread(
                 new_centroid_count[c_id] += 1
             new_centroid_sums_queue.put(new_centroid_sums)
             new_centroid_counts_queue.put(new_centroid_count)
-            logger.debug(f"Finished thread: {thread_id}")
+            logger.debug("Finished thread: %d", thread_id)
 
         def update_centroids():
             import multiprocessing as mp
 
-            logger.info("Updating centroids based on assignments.")
-            logger.info(f"Using {threads} threads.")
+            logger.debug("Updating centroids based on assignments.")
+            logger.debug("Using %d threads.", threads)
             global cents_t, vectors_t, assignments_t, new_centroid_thread_sums, new_centroid_thread_counts
             cents_t = centroids
             vectors_t = vectors
@@ -687,7 +682,7 @@ def update_centroids():
                 )
                 workers[i].join()
 
-            logger.info("Finished all threads, aggregating partial results.")
+            logger.debug("Finished all threads, aggregating partial results.")
             new_centroids = []
             for c_id in range(partitions):
                 cent = []
@@ -703,7 +698,7 @@ def update_centroids():
 
         logger = setup(config, verbose)
         with tiledb.scope_ctx(ctx_or_config=config):
-            logger.info("Reading input vectors.")
+            logger.debug("Reading input vectors.")
             vectors = read_input_vectors(
                 source_uri=source_uri,
                 source_type=source_type,
@@ -715,15 +710,15 @@ def update_centroids():
                 verbose=verbose,
                 trace_id=trace_id,
             )
-            logger.debug(f"Input centroids: {centroids[0:5]}")
-            logger.info("Assigning vectors to centroids")
+            logger.debug("Input centroids: %s", centroids[0:5])
+            logger.debug("Assigning vectors to centroids")
             km = KMeans()
             km._n_threads = threads
             km.cluster_centers_ = centroids
             assignments = km.predict(vectors)
-            logger.debug(f"Assignments: {assignments[0:100]}")
+            logger.debug("Assignments: %s", assignments[0:100])
             partial_new_centroids = update_centroids()
-            logger.debug(f"New centroids: {partial_new_centroids[0:5]}")
+            logger.debug("New centroids: %s", partial_new_centroids[0:5])
             return partial_new_centroids
 
     def compute_new_centroids(*argv):
@@ -753,7 +748,7 @@ def ingest_flat(
             group = tiledb.Group(array_uri)
             parts_array_uri = group[PARTS_ARRAY_NAME].uri
             target = tiledb.open(parts_array_uri, mode="w")
-            logger.info(f"Input vectors start_pos: {start}, end_pos: {end}")
+            logger.debug("Input vectors start_pos: %d, end_pos: %d", start, end)
 
             for part in range(start, end, batch):
                 part_end = part + batch
@@ -771,8 +766,8 @@ def ingest_flat(
                     trace_id=trace_id,
                 )
 
-                logger.debug(f"Vector read:{len(in_vectors)}")
-                logger.info(f"Writing data to array {parts_array_uri}")
+                logger.debug("Vector read: %d", len(in_vectors))
+                logger.debug("Writing data to array %s", parts_array_uri)
                 target[0:dimensions, start:end] = np.transpose(in_vectors)
             target.close()
 
@@ -789,7 +784,7 @@ def write_centroids(
             logger = setup(config, verbose)
             group = tiledb.Group(array_uri)
             centroids_uri = group[CENTROIDS_ARRAY_NAME].uri
-            logger.info(f"Writing centroids to array {centroids_uri}")
+            logger.debug("Writing centroids to array %s", centroids_uri)
             with tiledb.open(centroids_uri, mode="w") as A:
                 A[0:dimensions, 0:partitions] = np.transpose(np.array(centroids))
 
@@ -836,7 +831,7 @@ def ingest_vectors_udf(
             partial_write_array_index_uri = (
                 partial_write_array_dir_uri + "/" + INDEX_ARRAY_NAME + "/" + part_name
             )
-            logger.info(f"Input vectors start_pos: {part}, end_pos: {part_end}")
+            logger.debug("Input vectors start_pos: %d, end_pos: %d", part, part_end)
             if source_type == "TILEDB_ARRAY":
                 logger.debug("Start indexing")
                 ivf_index_tdb(
@@ -903,15 +898,15 @@ def compute_partition_indexes_udf(
                         partition_sizes[i] += int(partial_index) - int(prev_index)
                         prev_index = partial_index
                         i += 1
-            logger.debug(f"Partition sizes: {partition_sizes}")
+            logger.debug("Partition sizes: %s", partition_sizes)
             i = 0
-            sum = 0
+            _sum = 0
             for partition_size in partition_sizes:
-                indexes[i] = sum
-                sum += partition_size
+                indexes[i] = _sum
+                _sum += partition_size
                 i += 1
-            indexes[i] = sum
-            logger.debug(f"Partition indexes: {indexes}")
+            indexes[i] = _sum
+            logger.debug("Partition indexes: %d", indexes)
             index_array = tiledb.open(index_array_uri, mode="w")
             index_array[:] = indexes
 
@@ -927,9 +922,7 @@ def consolidate_partition_udf(
     ):
         logger = setup(config, verbose)
         with tiledb.scope_ctx(ctx_or_config=config):
-            logger.info(
-                f"Consolidating partitions {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 = (
@@ -969,14 +962,12 @@ 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(
-                f"Partitions start: {partition_id_start} end: {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.info(f"Consolidating partitions start: {part} end: {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]:
@@ -988,21 +979,20 @@ def consolidate_partition_udf(
                     if start_pos != end_pos:
                         raise ValueError("Incorrect partition size.")
                     continue
-                logger.debug(f"Read slices: {read_slices}")
+                logger.debug("Read slices: %s", read_slices)
                 ids = partial_write_array_ids_array.multi_index[read_slices]["values"]
                 vectors = partial_write_array_parts_array.multi_index[:, read_slices][
                     "values"
                 ]
 
-                logger.debug(
-                    f"Ids shape {ids.shape}, expected size: {end_pos - start_pos} expected range:({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.")
 
-                logger.info(f"Writing data to array: {parts_array_uri}")
+                logger.info("Writing data to array: %s", parts_array_uri)
                 parts_array[:, start_pos:end_pos] = vectors
-                logger.info(f"Writing data to array: {ids_array_uri}")
+                logger.info("Writing data to array: %s", ids_array_uri)
                 ids_array[start_pos:end_pos] = ids
             parts_array.close()
             ids_array.close()
@@ -1321,7 +1311,7 @@ def consolidate_and_vacuum(
             size = in_size
         logger.debug("Input dataset size %d", size)
         logger.debug("Input dataset dimensions %d", dimensions)
-        logger.debug(f"Vector dimension type {vector_type}")
+        logger.debug("Vector dimension type %s", vector_type)
         if partitions == -1:
             partitions = int(math.sqrt(size))
         if training_sample_size == -1:
@@ -1331,9 +1321,9 @@ def consolidate_and_vacuum(
                 workers = 10
         else:
             workers = 1
-        logger.debug(f"Partitions {partitions}")
-        logger.debug(f"Training sample size {training_sample_size}")
-        logger.debug(f"Number of workers {workers}")
+        logger.debug("Partitions %d", partitions)
+        logger.debug("Training sample size %d", training_sample_size)
+        logger.debug("Number of workers %d", workers)
 
         if input_vectors_per_work_item == -1:
             input_vectors_per_work_item = VECTORS_PER_WORK_ITEM
@@ -1348,10 +1338,7 @@ 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(
@@ -1367,15 +1354,10 @@ 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(