Build an HNSW Vector Index from Scratch in Pure Python

from __future__ import annotations
import math
import random
from collections import defaultdict
from typing import List, Dict, Optional, Tuple, Callable

class HNSWNode:
    """A single node in the HNSW graph."""
    def __init__(self, node_id: int, vector: List[float]):
        self.id = node_id
        self.vector = vector
        # connections[layer] = list of neighbor node_ids at that layer
        self.connections: Dict[int, List[int]] = defaultdict(list)
        self.max_level: int = 0  # highest layer this node participates in

    def __repr__(self) -> str:
        return f"HNSWNode(id={self.id}, dim={len(self.vector)}, layers={self.max_level + 1})"


def euclidean_distance(a: List[float], b: List[float]) -> float:
    """L2 distance between two vectors."""
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))


def cosine_distance(a: List[float], b: List[float]) -> float:
    """1 - cosine similarity (returns angular distance)."""
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0 or norm_b == 0:
        return 1.0
    return 1.0 - (dot / (norm_a * norm_b))

class HNSW:
    def __init__(
        self,
        dim: int,
        max_layers: int = 12,
        M: int = 16,
        ef_construction: int = 200,
        ef_search: int = 50,
        distance_fn: Callable[[List[float], List[float]], float] = euclidean_distance,
    ):
        self.dim = dim
        self.max_layers = max_layers
        self.M = M                    # max connections per node per layer
        self.ef_construction = ef_construction  # beam width during insert
        self.ef_search = ef_search    # beam width during query
        self.distance_fn = distance_fn

        self.nodes: Dict[int, HNSWNode] = {}
        self.entry_point: Optional[int] = None
        self.size = 0

    def _random_level(self) -> int:
        """Draw a layer from the exponential distribution used by the original paper."""
        level = 0
        # ml_ = -ln(0.5) ≈ 0.693; we use the common 0.5 probability per layer
        while random.random() < 0.5 and level < self.max_layers - 1:
            level += 1
        return level

    def _greedy_search(self, query: List[float], entry_id: int, layer: int) -> int:
        """Navigate the graph at `layer` to find the closest node to `query`."""
        current_id = entry_id
        visited = set()

        while True:
            visited.add(current_id)
            current_node = self.nodes[current_id]
            neighbors = current_node.connections.get(layer, [])

            best_id = current_id
            best_dist = self.distance_fn(query, current_node.vector)

            for neighbor_id in neighbors:
                if neighbor_id in visited:
                    continue
                dist = self.distance_fn(query, self.nodes[neighbor_id].vector)
                if dist < best_dist:
                    best_dist = dist
                    best_id = neighbor_id

            if best_id == current_id:
                # No neighbor is closer — we've reached the local optimum
                break
            current_id = best_id

        return current_id

    def _search_layer(
        self,
        query: List[float],
        entry_id: int,
        layer: int,
        ef: int,
    ) -> List[Tuple[float, int]]:
        """
        Return the `ef` closest nodes to `query` at `layer`.
        Uses a candidate set and a visited set, similar to the original paper.
        """
        # candidates: min-heap by distance (simulated with a sorted list for clarity)
        candidates: List[Tuple[float, int]] = [(self.distance_fn(query, self.nodes[entry_id].vector), entry_id)]
        visited = {entry_id}
        results: List[Tuple[float, int]] = []

        while candidates:
            # Pop the closest candidate
            candidates.sort(key=lambda x: x[0])
            current_dist, current_id = candidates.pop(0)
            results.append((current_dist, current_id))

            current_node = self.nodes[current_id]
            neighbors = current_node.connections.get(layer, [])

            for neighbor_id in neighbors:
                if neighbor_id in visited:
                    continue
                visited.add(neighbor_id)
                dist = self.distance_fn(query, self.nodes[neighbor_id].vector)
                candidates.append((dist, neighbor_id))

            # Prune: keep only the top `ef` candidates
            if len(candidates) > ef:
                candidates = candidates[:ef]

        return results[:ef]

    def add(self, vector: List[float]) -> int:
        """Insert a vector into the index and return its assigned node ID."""
        if len(vector) != self.dim:
            raise ValueError(f"Expected dim={self.dim}, got {len(vector)}")

        node_id = self.size
        self.size += 1
        new_node = HNSWNode(node_id, vector)
        self.nodes[node_id] = new_node

        # First node becomes the entry point
        if self.entry_point is None:
            self.entry_point = node_id
            new_node.max_level = 0
            return node_id

        # Step 1: Draw a random level for this node
        level = self._random_level()
        new_node.max_level = level

        # Step 2: Find the closest node at each layer, from top down
        entry = self.entry_point
        for layer in range(self.max_layers - 1, level - 1, -1):
            entry = self._greedy_search(vector, entry, layer)

        # Step 3: At each layer from `level` down to 0, connect to M nearest neighbors
        for layer in range(level, -1, -1):
            # Get ef_construction closest candidates at this layer
            candidates = self._search_layer(vector, entry, layer, self.ef_construction)

            # Select up to M neighbors to connect to
            selected = [nid for _, nid in candidates if nid != node_id][:self.M]

            # Add bidirectional connections
            new_node.connections[layer].extend(selected)
            for neighbor_id in selected:
                neighbor = self.nodes[neighbor_id]
                # Add reverse link
                if node_id not in neighbor.connections[layer]:
                    if len(neighbor.connections[layer]) >= self.M:
                        # Evict the farthest connection if at capacity
                        self._evict_farthest(neighbor, layer, vector)
                    neighbor.connections[layer].append(node_id)

            # The entry point for the next lower layer
            if selected:
                entry = selected[0]

        # Update global entry point if this node is on a higher layer
        if level > self.nodes[self.entry_point].max_level:
            self.entry_point = node_id

        return node_id

    def _evict_farthest(self, node: HNSWNode, layer: int, query: List[float]) -> None:
        """Replace the farthest neighbor with the new node if it's closer."""
        connections = node.connections[layer]
        farthest_id = max(connections, key=lambda nid: self.distance_fn(query, self.nodes[nid].vector))
        farthest_dist = self.distance_fn(query, self.nodes[farthest_id].vector)
        new_dist = self.distance_fn(query, node.vector)
        if new_dist < farthest_dist:
            connections.remove(farthest_id)
            connections.append(node.id)

    def search(self, query: List[float], k: int = 10) -> List[Tuple[float, int]]:
        """
        Return the `k` closest vectors to `query` as (distance, node_id) tuples.
        """
        if self.entry_point is None:
            return []

        # Phase 1: Descend from top layer to layer 0
        entry = self.entry_point
        for layer in range(self.max_layers - 1, -1, -1):
            entry = self._greedy_search(query, entry, layer)

        # Phase 2: Search at layer 0 with ef_search candidates
        candidates = self._search_layer(query, entry, 0, self.ef_search)

        # Return the top-k closest
        candidates.sort(key=lambda x: x[0])
        return candidates[:k]

    def save(self, path: str) -> None:
        """Save the index to a JSON file."""
        import json
        data = {
            "dim": self.dim,
            "max_layers": self.max_layers,
            "M": self.M,
            "ef_construction": self.ef_construction,
            "ef_search": self.ef_search,
            "nodes": {
                str(nid): {
                    "vector": node.vector,
                    "connections": {str(l): ids for l, ids in node.connections.items()},
                    "max_level": node.max_level,
                }
                for nid, node in self.nodes.items()
            },
            "entry_point": self.entry_point,
            "size": self.size,
        }
        with open(path, "w") as f:
            json.dump(data, f)

    @classmethod
    def load(cls, path: str) -> HNSW:
        """Load an index from a JSON file."""
        import json
        with open(path) as f:
            data = json.load(f)
        index = cls(
            dim=data["dim"],
            max_layers=data["max_layers"],
            M=data["M"],
            ef_construction=data["ef_construction"],
            ef_search=data["ef_search"],
        )
        for nid_str, ndata in data["nodes"].items():
            node = HNSWNode(int(nid_str), ndata["vector"])
            node.connections = defaultdict(list, {int(l): ids for l, ids in ndata["connections"].items()})
            node.max_level = ndata["max_level"]
            index.nodes[int(nid_str)] = node
        index.entry_point = data["entry_point"]
        index.size = data["size"]
        return index

python3 -c "
from hnsw import HNSW, euclidean_distance
import random

# --- Build a small index ---
random.seed(42)
dim = 128
index = HNSW(dim=dim, max_layers=10, M=16, ef_construction=100, ef_search=50)

# Insert 1000 random vectors
vectors = [[random.random() for _ in range(dim)] for _ in range(1000)]
for v in vectors:
    index.add(v)

print(f'Indexed {index.size} vectors')

# --- Query: find the 5 nearest neighbors ---
query = vectors[0]  # search for the vector itself
results = index.search(query, k=5)

print('\\nTop-5 results (should include the query vector itself at distance ~0):')
for dist, nid in results:
    print(f'  node {nid}: distance={dist:.6f}')

# --- Correctness check ---
# Brute-force the top-5 to compare
brute = sorted(
    [(euclidean_distance(query, v), i) for i, v in enumerate(vectors)],
    key=lambda x: x[0]
)[:5]
print('\\nBrute-force top-5:')
for dist, nid in brute:
    print(f'  node {nid}: distance={dist:.6f}')

# Recall check: does HNSW recall all brute-force top-5 IDs?
hnsw_ids = {nid for _, nid in results}
bf_ids = {nid for _, nid in brute}
recall = len(hnsw_ids & bf_ids) / len(bf_ids)
print(f'\\nRecall@5: {recall:.2%}')

# --- Persistence test ---
index.save('/tmp/test_hnsw.json')
loaded = HNSW.load('/tmp/test_hnsw.json')
print(f'\\nLoaded index: {loaded.size} nodes, entry_point={loaded.entry_point}')
"