Add support for randomly sampling the input data to compute centroids with (#183)

Author: jparismorgan

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

@@ -1,14 +1,22 @@
 import json
 from functools import partial
 from typing import Any, Mapping, Optional, Tuple
-
+import enum
+import random
 import numpy as np
-from tiledb.cloud.dag import Mode
 
+from tiledb.cloud.dag import Mode
 from tiledb.vector_search._tiledbvspy import *
 from tiledb.vector_search.storage_formats import STORAGE_VERSION, validate_storage_version
 
 
+class TrainingSamplingPolicy(enum.Enum):
+    FIRST_N = 1
+    RANDOM = 2
+
+    def __str__(self):
+        return self.name.replace("_", " ").title()
+
 def ingest(
     index_type: str,
     index_uri: str,
@@ -25,6 +33,7 @@ def ingest(
     namespace: Optional[str] = None,
     size: int = -1,
     partitions: int = -1,
+    training_sampling_policy: TrainingSamplingPolicy = TrainingSamplingPolicy.FIRST_N,
     copy_centroids_uri: str = None,
     training_sample_size: int = -1,
     training_input_vectors: np.ndarray = None,
@@ -33,6 +42,8 @@ def ingest(
     workers: int = -1,
     input_vectors_per_work_item: int = -1,
     max_tasks_per_stage: int= -1,
+    input_vectors_per_work_item_during_sampling: int = -1,
+    max_sampling_tasks: int= -1,
     storage_version: str = STORAGE_VERSION,
     verbose: bool = False,
     trace_id: Optional[str] = None,
@@ -102,6 +113,14 @@ def ingest(
     max_tasks_per_stage: int = -1
         Max number of tasks per execution stage of ingestion,
         if not provided, is auto-configured
+    input_vectors_per_work_item_during_sampling: int = -1
+        number of vectors per sample ingestion work item,
+        if not provided, is auto-configured
+        only valid with training_sampling_policy=TrainingSamplingPolicy.RANDOM
+    max_sampling_tasks: int = -1
+        Max number of tasks per execution stage of sampling,
+        if not provided, is auto-configured
+        only valid with training_sampling_policy=TrainingSamplingPolicy.RANDOM
     storage_version: str
         Vector index storage format version. If not provided, defaults to the latest version.
     verbose: bool
@@ -169,6 +188,10 @@ def ingest(
     for variable in ["copy_centroids_uri", "training_input_vectors", "training_source_uri", "training_source_type"]:
         if index_type != "IVF_FLAT" and locals().get(variable) is not None:
             raise ValueError(f"{variable} should only be provided with index_type IVF_FLAT")
+        
+    for variable in ["copy_centroids_uri", "training_input_vectors", "training_source_uri", "training_source_type"]:
+        if training_sampling_policy != TrainingSamplingPolicy.FIRST_N and locals().get(variable) is not None:
+            raise ValueError(f"{variable} should not provided alonside training_sampling_policy")
 
     # use index_group_uri for internal clarity
     index_group_uri = index_uri
@@ -191,6 +214,7 @@ def ingest(
     ]
     DEFAULT_ATTR_FILTERS = storage_formats[storage_version]["DEFAULT_ATTR_FILTERS"]
     VECTORS_PER_WORK_ITEM = 20000000
+    VECTORS_PER_SAMPLE_WORK_ITEM=1000000
     MAX_TASKS_PER_STAGE = 100
     CENTRALISED_KMEANS_MAX_SAMPLE_SIZE = 1000000
     DEFAULT_IMG_NAME = "3.9-vectorsearch"
@@ -302,9 +326,8 @@ def read_source_metadata(
         else:
             raise ValueError(f"Not supported source_type {source_type} - valid types are [TILEDB_ARRAY, U8BIN, F32BIN, FVEC, IVEC, BVEC]")
 
-    def write_input_vectors(
+    def create_array(
         group: tiledb.Group,
-        input_vectors: np.ndarray,
         size: int,
         dimensions: int,
         vector_type: np.dtype,
@@ -350,6 +373,18 @@ def write_input_vectors(
         tiledb.Array.create(input_vectors_array_uri, input_vectors_array_schema)
         group.add(input_vectors_array_uri, name=array_name)
 
+        return input_vectors_array_uri
+
+    def write_input_vectors(
+        group: tiledb.Group,
+        input_vectors: np.ndarray,
+        size: int,
+        dimensions: int,
+        vector_type: np.dtype,
+        array_name: str
+    ) -> str:
+        input_vectors_array_uri = create_array(group=group, size=size, dimensions=dimensions, vector_type=vector_type, array_name=array_name)
+
         input_vectors_array = tiledb.open(
             input_vectors_array_uri, "w", timestamp=index_timestamp
         )
@@ -797,6 +832,91 @@ def copy_centroids(
     # --------------------------------------------------------------------
     # centralised kmeans UDFs
     # --------------------------------------------------------------------
+    def random_sample_from_input_vectors(
+        source_uri: str,
+        source_type: str,
+        vector_type: np.dtype,
+        dimensions: int,
+        source_start_pos: int,
+        source_end_pos: int,
+        batch: int,
+        random_sample_size: int,
+        output_source_uri: str,
+        output_start_pos: int,
+        config: Optional[Mapping[str, Any]] = None,
+        verbose: bool = False,
+    ):
+        '''
+        Reads a random sample of vectors from the source data and appends them to the output array.
+
+        Parameters
+        ----------
+        source_uri: str
+            Data source URI.
+        source_type: str
+            Type of the source data.
+        vector_type: np.dtype
+            Type of the vectors.
+        dimensions: int
+            Number of dimensions in a vector.
+        vector_start_pos: int
+            Start position of source_uri to read from.
+        vector_end_pos: int
+            End position of source_uri to read to.
+        batch: int
+            Read the source data in batches of this size.
+        random_sample_size: int
+            Number of vectors to randomly sample from the source data.
+        output_source_uri: str
+            URI of the output array.
+        output_start_pos: int
+            Start position of the output array to write to.
+        '''
+        if random_sample_size == 0:
+            return
+
+        with tiledb.scope_ctx(ctx_or_config=config):
+            source_size = source_end_pos - source_start_pos
+            num_sampled = 0
+            for start in range(source_start_pos, source_end_pos, batch):
+                # What vectors to read from the source_uri.
+                end = start + batch
+                if end > source_end_pos:
+                    end = source_end_pos
+
+                # How many vectors sample from the vectors read.
+                percent_of_data_to_read = (end - start) / source_size
+                num_to_sample = math.ceil(random_sample_size * percent_of_data_to_read)
+                if num_sampled + num_to_sample > random_sample_size:
+                    num_to_sample = random_sample_size - num_sampled
+                if num_to_sample == 0:
+                    continue
+                num_sampled += num_to_sample
+
+                # Read from the source data.
+                vectors = read_input_vectors(
+                    source_uri=source_uri,
+                    source_type=source_type,
+                    vector_type=vector_type,
+                    dimensions=dimensions,
+                    start_pos=start,
+                    end_pos=end,
+                    config=config,
+                    verbose=verbose,
+                    trace_id=trace_id,
+                )
+
+                # Randomly sample from the data we read.
+                row_indices = np.random.choice(vectors.shape[0], size=num_to_sample, replace=False)
+                sampled_vectors = vectors[row_indices]
+
+                # Append to output array.
+                with tiledb.open(output_source_uri, mode="w", timestamp=index_timestamp) as A:
+                    A[0:dimensions, output_start_pos:output_start_pos + num_to_sample] = np.transpose(sampled_vectors)
+        
+        if num_sampled != random_sample_size:
+            raise ValueError(f"The random sampling within a batch ran into an issue: num_sampled ({num_sampled}) != random_sample_size ({random_sample_size})")
+
     def centralised_kmeans(
         index_group_uri: str,
         source_uri: str,
@@ -1564,6 +1684,8 @@ def create_ingestion_dag(
         training_source_type: Optional[str],
         input_vectors_per_work_item: int,
         input_vectors_work_items_per_worker: int,
+        input_vectors_per_work_item_during_sampling: int,
+        input_vectors_work_items_per_worker_during_sampling: int,
         table_partitions_per_work_item: int,
         table_partitions_work_items_per_worker: int,
         workers: int,
@@ -1600,6 +1722,15 @@ def create_ingestion_dag(
         input_vectors_batch_size = (
             input_vectors_per_work_item * input_vectors_work_items_per_worker
         )
+
+        # The number of vectors each task will read.
+        input_vectors_batch_size_during_sampling = (
+            # The number of vectors to read into memory in one batch within a task.
+            input_vectors_per_work_item_during_sampling * 
+            # The number of batches that a single task will need to run.
+            input_vectors_work_items_per_worker_during_sampling
+        )
+
         if index_type == "FLAT":
             ingest_node = submit(
                 ingest_flat,
@@ -1637,6 +1768,59 @@ def create_ingestion_dag(
                     image_name=DEFAULT_IMG_NAME,
                 )
             else:
+                random_sample_nodes = []
+                if training_sampling_policy == TrainingSamplingPolicy.RANDOM:
+                    # Create an empty array to write the sampled vectors to.
+                    group = tiledb.Group(index_group_uri, "w")
+                    training_source_uri = create_array(
+                        group=group,
+                        size=training_sample_size,
+                        dimensions=dimensions,
+                        vector_type=vector_type,
+                        array_name=TRAINING_INPUT_VECTORS_ARRAY_NAME
+                    )
+                    training_source_type = "TILEDB_ARRAY"
+                    group.close()
+
+                    idx = 0
+                    num_sampled = 0
+                    for start in range(0, in_size, input_vectors_batch_size_during_sampling):
+                        # What vectors to read from the source_uri.
+                        end = start + input_vectors_batch_size_during_sampling
+                        if end > size:
+                            end = size
+
+                        # How many vectors to sample from the vectors read.
+                        percent_of_data_to_read = (end - start) / in_size
+                        num_to_sample = math.ceil(training_sample_size * percent_of_data_to_read)
+                        if num_sampled + num_to_sample > training_sample_size:
+                            num_to_sample = training_sample_size - num_sampled
+                        if num_to_sample == 0:
+                            continue
+
+                        random_sample_nodes.append(submit(
+                            random_sample_from_input_vectors,
+                            source_uri=source_uri,
+                            source_type=source_type,
+                            vector_type=vector_type,
+                            dimensions=dimensions,
+                            source_start_pos=start,
+                            source_end_pos=end,
+                            batch=input_vectors_per_work_item_during_sampling,
+                            random_sample_size=num_to_sample,
+                            output_source_uri=training_source_uri,
+                            output_start_pos=num_sampled,
+                            config=config,
+                            verbose=verbose,
+                            name="read-random-sample-" + str(idx),
+                            resources={"cpu": str(threads), "memory": "1Gi"},
+                            image_name=DEFAULT_IMG_NAME,
+                        ))
+                        num_sampled += num_to_sample
+                        idx += 1
+                    if num_sampled != training_sample_size:
+                        raise ValueError(f"The random sampling ran into an issue: num_sampled ({num_sampled}) != training_sample_size ({training_sample_size})")
+
                 if training_sample_size <= CENTRALISED_KMEANS_MAX_SAMPLE_SIZE:
                     centroids_node = submit(
                         centralised_kmeans,
@@ -1657,6 +1841,9 @@ def create_ingestion_dag(
                         resources={"cpu": "8", "memory": "32Gi"},
                         image_name=DEFAULT_IMG_NAME,
                     )
+
+                    for random_sample_node in random_sample_nodes:
+                        centroids_node.depends_on(random_sample_node)
                 else:
                     internal_centroids_node = submit(
                         init_centroids,
@@ -1950,6 +2137,9 @@ def consolidate_and_vacuum(
         logger.debug("Input dataset size %d", size)
         logger.debug("Input dataset dimensions %d", dimensions)
         logger.debug("Vector dimension type %s", vector_type)
+        if training_sample_size > in_size:
+            raise ValueError(f"training_sample_size {training_sample_size} is larger than the input dataset size {in_size}")
+
         if partitions == -1:
             partitions = max(1, int(math.sqrt(size)))
         if training_sample_size == -1:
@@ -1976,6 +2166,7 @@ def consolidate_and_vacuum(
             if external_ids_type is None:
                 external_ids_type = "U64BIN"
 
+        # Compute task parameters for main ingestion.
         if input_vectors_per_work_item == -1:
             input_vectors_per_work_item = VECTORS_PER_WORK_ITEM
         input_vectors_work_items = int(math.ceil(size / input_vectors_per_work_item))
@@ -1996,6 +2187,27 @@ def consolidate_and_vacuum(
             input_vectors_work_items_per_worker,
         )
 
+        # Compute task parameters for random sampling.
+        # How many input vectors to read into memory in one batch within a task.
+        if input_vectors_per_work_item_during_sampling == -1:
+            input_vectors_per_work_item_during_sampling = VECTORS_PER_SAMPLE_WORK_ITEM
+        # How many total batches we need to read all the data..
+        input_vectors_work_items_during_sampling = int(math.ceil(size / input_vectors_per_work_item_during_sampling))
+        # The number of tasks to create, at max.
+        if max_sampling_tasks == -1:
+            max_sampling_tasks = MAX_TASKS_PER_STAGE
+        # The number of batches a single task will run. If there are more batches required than 
+        # allowed tasks, each task will process mutiple batches.
+        input_vectors_work_items_per_worker_during_sampling = 1
+        if input_vectors_work_items_during_sampling > max_sampling_tasks:
+            input_vectors_work_items_per_worker_during_sampling = int(
+                math.ceil(input_vectors_work_items_during_sampling / max_sampling_tasks)
+            )
+            input_vectors_work_items_during_sampling = max_sampling_tasks
+        logger.debug("input_vectors_per_work_item_during_sampling %d", input_vectors_per_work_item_during_sampling)
+        logger.debug("input_vectors_work_items_during_sampling %d", input_vectors_work_items_during_sampling)
+        logger.debug("input_vectors_work_items_per_worker_during_sampling %d", input_vectors_work_items_per_worker_during_sampling)
+
         vectors_per_table_partitions = max(1, size / partitions)
         table_partitions_per_work_item = max(
             1,
@@ -2054,6 +2266,8 @@ def consolidate_and_vacuum(
             training_source_type=training_source_type,
             input_vectors_per_work_item=input_vectors_per_work_item,
             input_vectors_work_items_per_worker=input_vectors_work_items_per_worker,
+            input_vectors_per_work_item_during_sampling=input_vectors_per_work_item_during_sampling,
+            input_vectors_work_items_per_worker_during_sampling=input_vectors_work_items_per_worker_during_sampling,
             table_partitions_per_work_item=table_partitions_per_work_item,
             table_partitions_work_items_per_worker=table_partitions_work_items_per_worker,
             workers=workers,

apis/python/test/common.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 import tiledb
-
+from tiledb.vector_search.storage_formats import storage_formats, STORAGE_VERSION
 
 def xbin_mmap(fname, dtype):
     n, d = map(int, np.fromfile(fname, dtype="uint32", count=2))
@@ -302,3 +302,11 @@ def check_equals(result_d, result_i, expected_result_d, expected_result_i):
 def random_name(name: str) -> str:
     suffix = "".join(random.choices(string.ascii_letters, k=10))
     return f"zzz_unittest_{name}_{suffix}"
+
+def check_training_input_vectors(index_uri: str, expected_training_sample_size: int, expected_dimensions: int):
+    training_input_vectors_uri = f"{index_uri}/{storage_formats[STORAGE_VERSION]['TRAINING_INPUT_VECTORS_ARRAY_NAME']}"
+    with tiledb.open(training_input_vectors_uri, mode="r") as src_array:
+        training_input_vectors = np.transpose(src_array[:, :]["values"])
+        assert training_input_vectors.shape[0] == expected_training_sample_size
+        assert training_input_vectors.shape[1] == expected_dimensions
+        assert not np.isnan(training_input_vectors).any()