Python API

Quake Index and Configuration

class quake.QuakeIndex

add(self: quake._bindings.QuakeIndex, arg0: torch.Tensor, arg1: torch.Tensor) → ModifyTimingInfo

Add new vectors to the index.

Parameters:

• x (Tensor) – Tensor of vectors to add.

• ids (Tensor) – Tensor of corresponding IDs.

build(self: quake._bindings.QuakeIndex, arg0: torch.Tensor, arg1: torch.Tensor, arg2: IndexBuildParams) → BuildTimingInfo

Build the index from a tensor of vectors and a corresponding tensor of IDs.

Parameters:

• x (Tensor) – Tensor of shape [num_vectors, dimension].

• ids (Tensor) – Tensor of shape [num_vectors].

• build_params (IndexBuildParams) – Parameters for building the index. Defaults include: - niter: default = 5 - metric: default = l2 - num_workers: default = 0

property current_level

The current level of the index.

get(self: quake._bindings.QuakeIndex, arg0: torch.Tensor) → torch.Tensor

Retrieve vectors from the index by ID.

Parameters:

ids (Tensor) – Tensor of IDs to retrieve.

get_ids(self: quake._bindings.QuakeIndex) → torch.Tensor

Return all vector IDs stored in the index.

initialize_maintenance_policy(self: quake._bindings.QuakeIndex, arg0: MaintenancePolicyParams) → None

Initialize the maintenance policy for the index.

Parameters:

maintenance_policy_params (MaintenancePolicyParams) – Parameters for the maintenance policy.

load(self: quake._bindings.QuakeIndex, arg0: str, arg1: int) → None

Load an index from a specified path.

Parameters:

• path (str) – The path from which to load the index.

• n_workers (int, optional) – Number of workers for query processing (default = 0).

maintenance(self: quake._bindings.QuakeIndex) → MaintenanceTimingInfo

Perform maintenance operations on the index (e.g., splits and merges). Returns timing information for the maintenance operation.

nlist(self: quake._bindings.QuakeIndex) → int

Return the number of partitions (lists) in the index.

ntotal(self: quake._bindings.QuakeIndex) → int

Return the total number of vectors stored in the index.

property parent

Return the parent index over the centroids.

remove(self: quake._bindings.QuakeIndex, arg0: torch.Tensor) → ModifyTimingInfo

Remove vectors from the index.

Parameters:

ids (Tensor) – Tensor of IDs to remove.

save(self: quake._bindings.QuakeIndex, arg0: str) → None

Save the index to a specified path.

Parameters:

path (str) – The path to save the index.

search(self: quake._bindings.QuakeIndex, arg0: torch.Tensor, arg1: SearchParams) → SearchResult

Search the index for nearest neighbors.

Parameters:

• x (Tensor) – Query tensor of shape [num_queries, dimension].

• search_params (SearchParams) – Parameters for the search operation. Defaults include: - k: default = 1 - nprobe: default = 1 - recall_target: default = -1.000000

class quake.SearchParams

property aps_flush_period_us

APS flush period in microseconds. default = 100

property batched_scan

Flag for batched scanning. default = 0

property initial_search_fraction

Initial fraction of partitions to search. default = 0.200000

property k

Number of neighbors to return. default = 1

property nprobe

Number of partitions to probe. default = 1

property num_threads

Number of threads to use for search within a single worker.

property recall_target

Recall target. default = -1.000000

property recompute_threshold

Threshold to trigger recomputation of APS. default = 0.010000

property use_precomputed

Flag to use precomputed inc beta fn for APS. default = 0

class quake.IndexBuildParams

property metric

Distance metric. default = l2

property niter

Number of k-means iterations. default = 5

property nlist

Number of clusters (lists). default = 0

property num_workers

Number of workers. default = 0

class quake.MaintenancePolicyParams

property alpha

Alpha parameter. default = 0.900000

property delete_threshold_ns

Delete threshold (ns). default = 20.000000

property enable_delete_rejection

Enable delete rejection. default = 1

property enable_split_rejection

Enable split rejection. default = 1

property k_large

Large k value for maintenance. default = 50

property k_small

Small k value for maintenance. default = 50

property maintenance_policy

Maintenance policy type. default = query_cost

property max_partition_ratio

Maximum allowed partition ratio. default = 2.000000

property min_partition_size

Minimum allowed partition size. default = 32

property modify_centroids

Flag to modify centroids during maintenance. default = 1

property refinement_iterations

Number of refinement iterations. default = 3

property refinement_radius

Radius for local partition refinement. default = 100

property split_threshold_ns

Split threshold (ns). default = 20.000000

property target_partition_size

Target partition size. default = 1000

property window_size

Window size for measuring hit rates. default = 1000

class quake.SearchResult

property distances

Distances to the nearest neighbors.

property ids

Indices of the nearest neighbors.

property timing_info

Timing information for the search operation.

class quake.MaintenanceTimingInfo

property delete_refine_time_us

Time taken for refinement of delete operations in microseconds.

property delete_time_us

Time taken for delete operations in microseconds.

property n_deletes

Number of partition delete operations performed.

property n_splits

Number of partition split operations performed.

property split_refine_time_us

Time taken for refinement of split operations in microseconds.

property split_time_us

Time taken for split operations in microseconds.

property total_time_us

Total time taken for maintenance in microseconds.

class quake.SearchTimingInfo

property boundary_distance_time_ns

Time spent on computing boundary distances in nanoseconds.

property buffer_init_time_ns

Time spent on initializing buffers in nanoseconds.

property job_enqueue_time_ns

Time spent on creating jobs in nanoseconds.

property job_wait_time_ns

Time spent waiting for jobs to complete in nanoseconds.

property n_clusters

Number of clusters searched.

property n_queries

Number of queries performed.

property parent_info

Search info for the parent index.

property partitions_scanned

Number of partitions scanned.

property result_aggregate_time_ns

Time spent on aggregating results in nanoseconds.

property search_params

Parameters used for the search operation.

property total_time_ns

Total time taken for the search operation in nanoseconds.

class quake.BuildTimingInfo

property assign_time_us

Time taken for assignment in microseconds.

property code_size

Size of PQ codes in bytes.

property d

Dimension of the vectors.

property n_codebooks

Number of codebooks in the index.

property n_vectors

Number of vectors in the index.

property total_time_us

Total time taken for the build operation in microseconds.

property train_time_us

Time taken for training in microseconds.

class quake.ModifyTimingInfo

property find_partition_time_us

Time taken to find the partition for the modify operation in microseconds.

property input_validation_time_us

Time taken for validation in microseconds.

property modify_count

property modify_time_us

Total time taken for the modify operation in microseconds.

Workload Generation and Evaluation

class quake.workload_generator.DynamicWorkloadGenerator(workload_dir: str | Path, base_vectors: ndarray, metric: str, insert_ratio: float, delete_ratio: float, query_ratio: float, update_batch_size: int, query_batch_size: int, number_of_operations: int, initial_size: int, cluster_size: int, cluster_sample_distribution: str, queries: ndarray, query_cluster_sample_distribution: str = 'uniform', seed: int = 1738, initial_clustering_path: str | Path | None = None, overwrite: bool = False)

Bases: object

Generates a dynamic workload from a static vector search dataset.

Workflow:

1. Cluster the base vectors.

2. Initialize the workload with an initial resident set.

3. Generate operations (insert, delete, query) according to given ratios.

4. Save each operation and a runbook that includes a summary.

generate_workload()

Generate the workload and print a summary at the end.

initialize_clustered_index()

initialize_workload()

sample(size: int, operation_type: str)

validate_parameters()

workload_exists()

class quake.workload_generator.StratifiedClusterSampler(assignments: Tensor, centroids: Tensor)

Bases: VectorSampler

Sample vectors from clusters.

This sampler uses cluster assignments and centroid distances to sample in a stratified fashion.

sample(sample_pool: Tensor, size: int, update_ranks: bool = True)

Sample vectors for an operation.

update_ranks(root_cluster: int)

class quake.workload_generator.UniformSampler

Bases: VectorSampler

Uniformly sample vectors.

sample(sample_pool: Tensor, size: int, update_ranks: bool = True)

Sample vectors for an operation.

class quake.workload_generator.VectorSampler

Bases: ABC

Abstract class for sampling vectors.

abstractmethod sample(size: int)

Sample vectors for an operation.

class quake.workload_generator.WorkloadEvaluator(workload_dir: str | Path, output_dir: str | Path, base_vectors_path: str | Path | None = None)

Bases: object

Evaluates a generated workload on a given index and produces summary statistics and plots.

evaluate_workload(name, index, build_params, search_params, do_maintenance=False)

Evaluate the workload on the index. At the end a summary is printed and a multi-panel plot is saved.

initialize_index(name, index, build_params)

quake.workload_generator.run_query(index, queries: Tensor, search_k: int, search_params: dict, gt_ids: Tensor, single_query: bool = True, tune_nprobe: bool = False, nprobes: Tensor | None = None)

Run queries on the index and compute the recall.

Built-in Datasets

class quake.datasets.ann_datasets.Dataset

Bases: ABC

abstractmethod download(overwrite: bool = False)

download_dir: Path = None

abstractmethod is_downloaded() → bool

load() → List[ndarray | Tensor]

abstractmethod load_ground_truth() → ndarray | Tensor

abstractmethod load_queries() → ndarray | Tensor

abstractmethod load_vectors() → ndarray | Tensor

metric: str = 'l2'

url: str = None

class quake.datasets.ann_datasets.Sift1m(download_dir: str | Path = PosixPath('data'))

Bases: Dataset

download(overwrite: bool = False)

is_downloaded() → bool

load_ground_truth() → ndarray | Tensor

load_queries() → ndarray | Tensor

load_vectors() → ndarray | Tensor

url: str = 'ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz'

quake.datasets.ann_datasets.load_dataset(name: str, download_dir: str = PosixPath('data'), overwrite_download: bool = False) → List[ndarray | Tensor]

Index Wrappers

class quake.index_wrappers.quake.QuakeWrapper

Bases: IndexWrapper

add(vectors: Tensor, ids: Tensor | None = None, num_threads: int = 0)

Add vectors to the index.

Parameters:

• vectors – The vectors to add to the index.

• ids – The ids of the vectors to add to the index.

assignments: Tensor | None

build(vectors: Tensor, nc: int, metric: str = 'l2', ids: Tensor | None = None, n_workers: int = 0, m: int = -1, code_size: int = 8)

Build the index with the given vectors and arguments.

Parameters:

• vectors – The vectors to build the index with.

• nc – The number of centroids (ivf).

• metric – The distance metric to use, optional. Default is “l2”.

• ids – The ids of the vectors.

centroids() → Tensor

Return the centroids of the index.

Returns:

The centroids of the index

cluster_ids() → Tensor

Return the cluster assignments of the vectors in the index.

Returns:

The cluster ids of the index

d() → int

Return the dimension of the vectors in the index.

Returns:

The dimension of the vectors in the index.

index: QuakeIndex

index_state() → dict

Return the state of the index.

• n_list: The number of centroids in the index.

• `n_total’: The number of vectors in the index.

• metric: The distance metric used in the index.

Returns:

The state of the index as a dictionary.

load(filename: str, n_workers: int = 1, use_numa: bool = False, verbose: bool = False, verify_numa: bool = False, same_core: bool = True, use_centroid_workers: bool = False, use_adaptive_n_probe: bool = False)

Load the index from a file.

Parameters:

filename – The name of the file to load the index from.

maintenance()

Perform maintenance on the index. :return: maintenance results

metric() → str

Return the metric of the index.

Returns:

The metric of the index.

n_total() → int

Return the number of vectors in the index.

Returns:

The number of vectors in the index.

remove(ids: Tensor)

Remove vectors from the index.

Parameters:

indices – The indices of the vectors to remove.

save(filename: str)

Save the index to a file.

Parameters:

filename – The name of the file to save the index to.

search(query: Tensor, k: int, nprobe: int = 1, batched_scan=False, recall_target: float = -1, k_factor=4.0, use_precomputed=True) → Tuple[Tensor, Tensor]

Find the k-nearest neighbors of the query vectors.

Parameters:

• query – The query vectors.

• k – The number of nearest neighbors to find.

• nprobe – The number of centroids to visit during search. Default is 1.

Returns:

The distances and indices of the k-nearest neighbors.

class quake.index_wrappers.faiss_ivf.FaissIVF

Bases: IndexWrapper

Wrapper class for faiss indexes.

add(vectors: Tensor, ids: Tensor | None = None, num_threads: int = 0)

Add vectors to the index.

Parameters:

• vectors – The vectors to add to the index.

• ids – The ids of the vectors to add to the index.

Raises:

AssertionError – If the vectors are not 2-dimensional.

build(vectors: Tensor, nc: int, m: int = 0, b: int = 0, metric: str = 'l2', ids: Tensor | None = None)

Build the index with the given vectors and arguments. When nc is 0, the index is a flat index or PQ index. When nc is not 0, the index is an IVF or IVFPQ index. (PQ is presence only if m and b are not 0)

Parameters:

• vectors – The vectors to build the index with.

• nc – The number of centroids (ivf).

• m – The number of subquantizers (pq), optional. Default is 0.

• b – The number of bytes per code (pq), optional. Default is 0.

• metric – The distance metric to use, optional. Default is “l2”.

Raises:

AssertionError – If any of the constraints on nc, m, or b are violated.

centroids() → Tensor | None

Return the centroids of the index if it is an IVF index. Otherwise, return None.

Returns:

The centroids of the index or None if the index is not an IVF index.

Return type:

Union[torch.Tensor, None]

d() → int

Return the dimension of the vectors in the index.

Returns:

The dimension of the vectors in the index.

index: Index

index_state() → dict

Return the state of the index.

• n_list: The number of centroids in the index.

• `n_total’: The number of vectors in the index.

• metric: The distance metric used in the index.

Returns:

The state of the index as a dictionary.

index_type: IndexType

load(filename: str)

Load the index from a file.

Parameters:

filename – The name of the file to load the index from.

maintenance()

Perform any necessary maintenance operations on the index

n_total() → int

Return the number of vectors in the index.

Returns:

The number of vectors in the index.

remove(ids: Tensor)

Remove vectors from the index. :param ids: The ids of the vectors to remove.

save(filename: str)

Save the index to a file.

Parameters:

filename – The name of the file to save the index to.

search(query: Tensor, k: int, nprobe: int = 1, rf: int = 1, batched_scan: bool = False) → Tuple[Tensor, Tensor]

Find the k-nearest neighbors of the query vectors.

Parameters:

• query – The query vectors.

• k – The number of nearest neighbors to find.

• nprobe – The number of centroids to visit during search. Default is 1. Only used for IVF indexes.

• rf – The re-ranking factor. Search for rf * k results to re-rank. Default is 1. Only used for PQ indexes.

Returns:

The distances and indices of the k-nearest neighbors.

class quake.index_wrappers.faiss_ivf.IndexType(value)

Bases: Enum

Enum class for the different types of indexes.

FLAT = 0

IVF = 2

IVFPQ = 3

PQ = 1

quake.index_wrappers.faiss_ivf.faiss_metric_to_str(metric: int) → str

Convert the faiss metric to the corresponding string. :param metric: The faiss metric to convert. :return: The string metric.

quake.index_wrappers.faiss_ivf.metric_str_to_faiss(metric: str) → int

Convert the metric to the corresponding faiss metric. :param metric: The metric to convert. :return: The faiss metric.