Vector Search：从暴力搜索到向量数据库Vector Search: From Brute Force to ANN Indexes

import heapq
import numpy as np

def top_k(query, documents, k=3, metric="cosine"):
    q = np.asarray(query, dtype=np.float32)
    d = np.asarray(documents, dtype=np.float32)
    if d.ndim != 2 or q.shape != (d.shape[1],):
        raise ValueError("query and document dimensions must match")
    if not 1 <= k <= len(d):
        raise ValueError("k must be between 1 and the document count")

    if metric == "cosine":
        q = q / np.maximum(np.linalg.norm(q), 1e-12)
        d = d / np.maximum(np.linalg.norm(d, axis=1, keepdims=True), 1e-12)
        values = d @ q                    # larger is better
        ranking_scores = values
    elif metric == "inner_product":
        values = d @ q                    # magnitude still matters
        ranking_scores = values
    elif metric == "l2":
        values = np.linalg.norm(d - q, axis=1)  # smaller is better
        ranking_scores = -values           # negate only for one max-heap API
    else:
        raise ValueError(metric)

    # O(N log K): never sort all N rows when K is small.
    top_ids = heapq.nlargest(
        k,
        range(len(ranking_scores)),
        key=lambda index: (float(ranking_scores[index]), -index),
    )
    return [(index, float(values[index])) for index in top_ids]

def main():
    labels = ["A", "B", "C", "D", "E"]
    query = np.array([1.00, 0.45], dtype=np.float32)
    documents = np.array([
        [0.72, 0.32],   # A: almost the same direction, shorter
        [1.25, 0.82],   # B: aligned and longer
        [1.04, 0.41],   # C: closest coordinate position
        [-0.12, 0.96],  # D: different direction
        [0.48, 0.78],   # E: medium direction and distance
    ], dtype=np.float32)

    print("query:", [round(float(value), 3) for value in query])
    for metric in ("cosine", "inner_product", "l2"):
        results = top_k(query, documents, k=3, metric=metric)
        rule = "larger is better" if metric != "l2" else "smaller is better"
        print(f"
{metric} Top-3 ({rule})")
        for rank, (index, value) in enumerate(results, start=1):
            print(f"  #{rank} document {labels[index]}: {value:.6f}")

if __name__ == "__main__":
    main()

import heapq
from time import perf_counter
import numpy as np

def exact_cosine_topk(matrix, query, k):
    # Offline contract: rows and query are already L2-normalized.
    heap = []
    for doc_id, vector in enumerate(matrix):          # N visits
        score = float(vector @ query)                 # d multiply-adds
        item = (score, doc_id)
        if len(heap) < k:
            heapq.heappush(heap, item)
        elif score > heap[0][0]:
            heapq.heapreplace(heap, item)             # log K
    return sorted(heap, reverse=True)

for n in [100, 1_000, 10_000, 100_000]:
    docs = np.random.default_rng(7).normal(size=(n, 32)).astype("float32")
    docs /= np.maximum(np.linalg.norm(docs, axis=1, keepdims=True), 1e-12)
    query = docs[0]
    started = perf_counter()
    result = exact_cosine_topk(docs, query, k=10)
    print(n, (perf_counter() - started) * 1000, "ms", result[:2])

from heapq import heappush, heappop
import numpy as np

def squared_l2(query, vector):
    delta = query - vector
    return float(delta @ delta)

def greedy_descent(query, entry_id, vectors, neighbors):
    """ef=1 search used on each sparse upper layer."""
    current = entry_id
    moves = []
    while True:
        current_distance = squared_l2(query, vectors[current])
        best_id, best_distance = current, current_distance
        for neighbor_id in neighbors.get(current, []):
            distance = squared_l2(query, vectors[neighbor_id])
            if (distance, neighbor_id) < (best_distance, best_id):
                best_id, best_distance = neighbor_id, distance
        if best_id == current:
            return current, moves
        moves.append((current, best_id, current_distance, best_distance))
        current = best_id

def search_layer(query, entry_ids, ef, vectors, neighbors):
    """Best-first layer search with candidates min-heap and results max-heap."""
    visited = set(entry_ids)
    candidates = []      # (distance, id): nearest expandable candidate first
    results = []         # (-distance, id): farthest retained result first

    for node_id in entry_ids:
        distance = squared_l2(query, vectors[node_id])
        heappush(candidates, (distance, node_id))
        # Negate both fields: heap root is the farthest node; for a distance
        # tie, the larger ID is considered worse and discarded first.
        heappush(results, (-distance, -node_id))

    while candidates:
        candidate_distance, candidate_id = heappop(candidates)
        worst_result_distance = -results[0][0]
        worst_result_id = -results[0][1]
        if len(results) >= ef and (candidate_distance, candidate_id) > (worst_result_distance, worst_result_id):
            break

        for neighbor_id in neighbors.get(candidate_id, []):
            if neighbor_id in visited:
                continue
            visited.add(neighbor_id)
            distance = squared_l2(query, vectors[neighbor_id])
            worst_result_distance = -results[0][0]
            worst_result_id = -results[0][1]
            if len(results) < ef or (distance, neighbor_id) < (worst_result_distance, worst_result_id):
                heappush(candidates, (distance, neighbor_id))
                heappush(results, (-distance, -neighbor_id))
                if len(results) > ef:
                    heappop(results)  # discard the farthest retained node

    ordered = sorted(
        (-negative_distance, -negative_id)
        for negative_distance, negative_id in results
    )
    return ordered, visited

def hnsw_search(query, vectors, layers, entry_id, ef_search, k):
    if len(vectors) == 0 or not layers:
        raise ValueError("vectors and layers must not be empty")
    if not 1 <= k <= len(vectors):
        raise ValueError("k must be between 1 and the vector count")
    if ef_search < k:
        raise ValueError("ef_search must be at least k")
    if not 0 <= entry_id < len(vectors):
        raise ValueError("entry_id is outside the vector table")

    current = entry_id
    descent_trace = []
    for layer in range(len(layers) - 1, 0, -1):
        current, moves = greedy_descent(query, current, vectors, layers[layer])
        descent_trace.extend((layer, *move) for move in moves)

    retained, visited = search_layer(
        query,
        entry_ids=[current],
        ef=ef_search,
        vectors=vectors,
        neighbors=layers[0],
    )
    return retained[:k], descent_trace, visited

def main():
    # A tiny hand-built 3-layer graph. Production HNSW builds these edges.
    vectors = np.array([
        [0.0, 0.00], [1.0, 0.15], [2.0, -0.05],
        [3.0, 0.10], [4.0, 0.00], [5.0, -0.10],
        [6.0, 0.05], [7.0, 0.12], [8.0, 0.00],
    ], dtype=np.float32)
    layer0 = {
        node: [neighbor for neighbor in (node - 1, node + 1) if 0 <= neighbor < len(vectors)]
        for node in range(len(vectors))
    }
    layer1 = {0: [2], 2: [0, 4], 4: [2, 6], 6: [4, 8], 8: [6]}
    layer2 = {0: [4], 4: [0, 8], 8: [4]}
    layers = [layer0, layer1, layer2]

    query = np.array([7.20, 0.05], dtype=np.float32)
    k, ef_search = 3, 5
    ann, descent_trace, visited = hnsw_search(
        query, vectors, layers, entry_id=0, ef_search=ef_search, k=k
    )
    exact = sorted(
        (squared_l2(query, vector), node_id)
        for node_id, vector in enumerate(vectors)
    )[:k]
    ann_ids = [node_id for _, node_id in ann]
    exact_ids = [node_id for _, node_id in exact]
    recall = len(set(ann_ids) & set(exact_ids)) / k

    print("query:", [round(float(value), 3) for value in query], "entry: v0")
    for layer, from_id, to_id, before, after in descent_trace:
        print(f"L{layer}: v{from_id} ({before:.4f}) -> v{to_id} ({after:.4f})")
    print("layer-0 visited:", sorted(visited))
    print("ANN Top-K:", [(f"v{node_id}", round(distance, 6)) for distance, node_id in ann])
    print("Exact Top-K:", [(f"v{node_id}", round(distance, 6)) for distance, node_id in exact])
    print(f"Recall@{k}: {recall:.1%}")

if __name__ == "__main__":
    main()

import numpy as np

def train_ivf(training_vectors, nlist, iterations=20, seed=7):
    # A compact, runnable L2 k-means trainer for teaching.
    if training_vectors.ndim != 2 or len(training_vectors) == 0:
        raise ValueError("training_vectors must be a non-empty 2-D array")
    if not 1 <= nlist <= len(training_vectors):
        raise ValueError("nlist must be between 1 and the training-vector count")
    if iterations < 1:
        raise ValueError("iterations must be positive")
    rng = np.random.default_rng(seed)
    initial_ids = rng.choice(len(training_vectors), nlist, replace=False)
    centroids = training_vectors[initial_ids].copy()

    for _ in range(iterations):
        distance = ((training_vectors[:, None, :] - centroids[None, :, :]) ** 2).sum(axis=2)
        assignments = distance.argmin(axis=1)
        updated = np.vstack([
            training_vectors[assignments == cell].mean(axis=0)
            if np.any(assignments == cell) else centroids[cell]
            for cell in range(nlist)
        ])
        converged = np.allclose(updated, centroids)
        centroids = updated
        if converged:
            break
    return centroids

def add_ivf(vectors, centroids):
    # Add is separate from training: assign the full corpus once.
    distance = ((vectors[:, None, :] - centroids[None, :, :]) ** 2).sum(axis=2)
    assignments = distance.argmin(axis=1)
    return [np.flatnonzero(assignments == cell) for cell in range(len(centroids))]

def sorted_smallest_ids(distance, count):
    if count < 0:
        raise ValueError("count must not be negative")
    count = min(count, len(distance))
    if count == 0:
        return np.empty(0, dtype=np.int64)

    # Partial selection is O(N); explicitly resolve the cutoff tie by ID so
    # repeated runs have the same result, then sort only the selected IDs.
    cutoff = np.partition(distance, count - 1)[count - 1]
    lower_ids = np.flatnonzero(distance < cutoff)
    tied_ids = np.flatnonzero(distance == cutoff)
    ids = np.concatenate([lower_ids, tied_ids[:count - len(lower_ids)]])
    return ids[np.lexsort((ids, distance[ids]))]

def search_ivf(query, vectors, centroids, lists, nprobe, k):
    if not 1 <= nprobe <= len(centroids):
        raise ValueError("nprobe must be between 1 and nlist")
    if not 1 <= k <= len(vectors):
        raise ValueError("k must be between 1 and the vector count")

    # Coarse stage: rank centroids and open the nearest nprobe lists.
    cell_distance = ((centroids - query) ** 2).sum(axis=1)
    selected = sorted_smallest_ids(cell_distance, nprobe)
    candidate_ids = np.concatenate([lists[cell] for cell in selected])
    if len(candidate_ids) == 0:
        raise ValueError("the selected IVF lists contain no vectors")

    # IVF-Flat fine stage: exact L2 only inside the candidate union.
    candidate_distance = ((vectors[candidate_ids] - query) ** 2).sum(axis=1)
    local_topk = sorted_smallest_ids(candidate_distance, k)
    return candidate_ids[local_topk]

def main():
    rng = np.random.default_rng(11)
    true_centers = np.array([
        [-1.0, -0.8],
        [0.0, 1.2],
        [1.1, -0.3],
    ], dtype=np.float32)

    # Training samples learn the coarse centroids; the corpus is added later.
    training_vectors = np.vstack([
        rng.normal(center, 0.16, size=(80, 2)) for center in true_centers
    ]).astype(np.float32)
    vectors = np.vstack([
        rng.normal(center, 0.18, size=(40, 2)) for center in true_centers
    ]).astype(np.float32)

    nlist = 3
    centroids = train_ivf(training_vectors, nlist=nlist, seed=7)
    lists = add_ivf(vectors, centroids)

    query = np.array([1.05, -0.35], dtype=np.float32)
    nprobe, k = 2, 5
    ann_ids = search_ivf(query, vectors, centroids, lists, nprobe, k)

    cell_distance = ((centroids - query) ** 2).sum(axis=1)
    selected_cells = sorted_smallest_ids(cell_distance, nprobe)
    candidate_count = sum(len(lists[cell]) for cell in selected_cells)
    exact_distance = ((vectors - query) ** 2).sum(axis=1)
    exact_ids = sorted_smallest_ids(exact_distance, k)
    recall = len(set(ann_ids.tolist()) & set(exact_ids.tolist())) / k

    print("trained centroids:
", np.round(centroids, 3))
    print("list sizes:", [len(postings) for postings in lists])
    print("query:", [round(float(value), 3) for value in query])
    print("selected cells:", selected_cells.tolist())
    print(f"candidate vectors: {candidate_count}/{len(vectors)}")
    print("ANN Top-K:", ann_ids.tolist())
    print("Exact Top-K:", exact_ids.tolist())
    print(f"Recall@{k}: {recall:.1%}")

if __name__ == "__main__":
    main()

from math import ceil
from statistics import mean

def recall_at_k(exact_ids, ann_ids, k):
    truth = set(exact_ids[:k])
    returned = set(ann_ids[:k])
    return len(truth & returned) / k

def nearest_rank_percentile(values, percentile):
    """Teaching definition: sorted_values[ceil(p/100 * n) - 1]."""
    if not values:
        raise ValueError("values must not be empty")
    if not 0 < percentile <= 100:
        raise ValueError("percentile must be in (0, 100]")
    ordered = sorted(values)
    rank = ceil(percentile / 100 * len(ordered))
    return ordered[rank - 1]

def main():
    # Five held-out queries. In production, exact_results come from a full scan
    # and ann_results come from the index under one fixed search configuration.
    exact_results = [list(range(start, start + 10)) for start in range(0, 100, 20)]
    ann_results = [
        exact_results[0],
        exact_results[1][:9] + [999],
        exact_results[2][:8] + [998, 999],
        exact_results[3][:7] + [997, 998, 999],
        exact_results[4][:6] + [996, 997, 998, 999],
    ]
    recalls = [
        recall_at_k(exact_ids, ann_ids, k=10)
        for exact_ids, ann_ids in zip(exact_results, ann_results)
    ]

    # 101 measured query latencies in milliseconds. Replace these samples with
    # perf_counter timings around the ANN search call on the target machine.
    latencies_ms = [6.0] * 51 + [10.0] * 45 + [40.0] * 4 + [120.0]

    print("per-query Recall@10:", recalls)
    print(f"mean Recall@10: {mean(recalls):.1%}")
    print(f"mean latency: {mean(latencies_ms):.2f} ms")
    for percentile in (50, 95, 99):
        value = nearest_rank_percentile(latencies_ms, percentile)
        print(f"p{percentile}: {value:.2f} ms")
    print(f"max: {max(latencies_ms):.2f} ms")

if __name__ == "__main__":
    main()