Vector Index Internals: Exact Search, IVF, HNSW, Product Quantization and Filtered Search with FAISS

"""Vector search for Kittiwake Outfitters. You write the functions marked TODO, one step at a time."""
import time

import faiss
import numpy as np

K = 10


# ---------- Step 1: exact search, and what "recall" means ----------

def configure():
    """Make faiss use one thread (the pod has one CPU but sees 48) and return faiss.omp_get_max_threads()."""
    # TODO (Step 1): faiss.omp_set_num_threads(1), then return faiss.omp_get_max_threads().
    raise NotImplementedError("Step 1: write configure()")


def exact_search(db, queries, k=K, batch=64):
    """(scores, ids), each (len(queries), k): the k products with the highest dot product per query, best first,
    computed with numpy in batches of `batch` queries (a full queries x products matrix would not fit in memory)."""
    # TODO (Step 1): for each batch of queries: s = batch @ db.T, np.argpartition(-s, k, axis=1)[:, :k],
    # then sort those k by score, descending; stack the batches.
    raise NotImplementedError("Step 1: write exact_search()")


def recall_at_k(found, truth):
    """The mean over queries of |found row & truth row| / k: the share of the true top k that the search returned."""
    # TODO (Step 1): per query, len(set(found without -1) & set(truth)) / k; the mean, as a float.
    raise NotImplementedError("Step 1: write recall_at_k()")


def ms_per_query(index, queries, k=K):
    """Milliseconds per query for index.search(queries, k), the faster of two runs."""
    best = None
    for _ in range(2):
        t = time.perf_counter()
        index.search(queries, k)
        dt = (time.perf_counter() - t) * 1000 / len(queries)
        best = dt if best is None else min(best, dt)
    return best


# ---------- Step 2: IVF, search a few clusters ----------

def build_ivf(db, nlist, train_size=20_000, seed=0):
    """An IndexIVFFlat with inner-product metric and nlist lists, trained on train_size random rows of db
    (np.random.default_rng(seed), without replacement), with every row of db added."""
    raise NotImplementedError("build_ivf() arrives in Step 2")


def set_param(index, name, value):
    """Set a search-time knob: "nprobe" on an IVF index, "efSearch" on an HNSW index."""
    if name == "efSearch":
        index.hnsw.efSearch = value
    else:
        setattr(faiss.extract_index_ivf(index) if name == "nprobe" else index, name, value)


def sweep(index, queries, truth, name, values, k=K):
    """For each value of the knob: {"value", "recall", "ms"} (recall_at_k against truth, ms_per_query)."""
    raise NotImplementedError("sweep() arrives in Step 2")


def cheapest(results, target):
    """The first result (in the order swept) whose recall reaches target, or None."""
    raise NotImplementedError("cheapest() arrives in Step 2")


# ---------- Step 3: HNSW, a graph to walk ----------

def build_hnsw(db, M, ef_construction):
    """An IndexHNSWFlat with inner-product metric, M links per node and efConstruction set, with db added."""
    raise NotImplementedError("build_hnsw() arrives in Step 3")


def index_mb(index):
    """The index's size in megabytes (1e6 bytes), as faiss.serialize_index() writes it."""
    raise NotImplementedError("index_mb() arrives in Step 3")


# ---------- Step 4: compress, then re-rank ----------

def build_ivfpq(db, nlist, m, train_size=20_000, seed=0):
    """An IndexIVFPQ (inner product, nlist lists, m sub-quantizers of 8 bits) trained like build_ivf(), db added."""
    raise NotImplementedError("build_ivfpq() arrives in Step 4")


def rerank(vectors, queries, candidates, k=K):
    """Exact re-scoring: for each query, the k candidate ids (ignore -1) with the highest dot product against the
    full vectors (which may be a memory-mapped array on disk), best first."""
    raise NotImplementedError("rerank() arrives in Step 4")


# ---------- Step 5: only what is in stock ----------

def post_filter(index, queries, allowed, k=K):
    """Search for k, then drop results that are not allowed (a boolean array over product ids). Shown for comparison."""
    _, ids = index.search(queries, k)
    return [[i for i in row if i >= 0 and allowed[i]] for row in ids]


def filtered_search(index, queries, allowed, k=K, nprobe=None):
    """Search only among allowed products: pass an IDSelectorBatch of the allowed ids in faiss.SearchParametersIVF
    (with nprobe if given) so the index skips everything else. Returns the ids, (len(queries), k)."""
    raise NotImplementedError("filtered_search() arrives in Step 5")


def exact_filtered(db, queries, allowed, k=K):
    """The true top k among allowed products only, as ids of the full catalogue."""
    raise NotImplementedError("exact_filtered() arrives in Step 5")