Building a Vector Search Engine from Scratch: The Math and Mechanics of HNSW Indexing

class VectorNode {
  public id: string;
  public vector: number[];
  public connections: Map<number, Set<VectorNode>>; // Layer level -> Connected nodes

  constructor(id: string, vector: number[]) {
    this.id = id;
    this.vector = vector;
    this.connections = new Map();
  }

  public addConnection(level: number, node: VectorNode) {
    if (!this.connections.has(level)) {
      this.connections.set(level, new Set());
    }
    this.connections.get(level)!.add(node);
  }
}

export class HNSWIndex {
  private maxLevel: number;
  private M: number; // Max connections per node
  private efConstruction: number; // Search depth during index build
  private enterNode: VectorNode | null = null;
  private levelMultiplier: number;

  constructor(M: number = 16, efConstruction: number = 200) {
    this.M = M;
    this.efConstruction = efConstruction;
    this.maxLevel = 0;
    this.levelMultiplier = 1 / Math.log(M);
  }

  // Calculate Cosine Similarity between two floating arrays
  private computeCosine(v1: number[], v2: number[]): number {
    let dot = 0;
    let mag1 = 0;
    let mag2 = 0;
    for (let i = 0; i < v1.length; i++) {
      dot += v1[i] * v2[i];
      mag1 += v1[i] * v1[i];
      mag2 += v2[i] * v2[i];
    }
    return dot / (Math.sqrt(mag1) * Math.sqrt(mag2));
  }

  // Insert a new Vector into the Multi-Layer graph skip list
  public insert(id: string, vector: number[]) {
    const newNode = new VectorNode(id, vector);
    
    // Determine random maximum layer depth using exponential decay distribution
    const nodeMaxLevel = Math.floor(-Math.log(Math.random()) * this.levelMultiplier);
    
    if (this.enterNode === null) {
      this.enterNode = newNode;
      this.maxLevel = nodeMaxLevel;
      return;
    }

    let currNode = this.enterNode;
    let queryDist = this.computeCosine(vector, currNode.vector);

    // 1. Search sparse highway layers down to nodeMaxLevel
    for (let level = this.maxLevel; level > nodeMaxLevel; level--) {
      let changed = true;
      while (changed) {
        changed = false;
        const neighbors = currNode.connections.get(level) || new Set();
        for (const neighbor of neighbors) {
          const dist = this.computeCosine(vector, neighbor.vector);
          if (dist > queryDist) {
            queryDist = dist;
            currNode = neighbor;
            changed = true;
          }
        }
      }
    }

    // 2. Insert node into layers from nodeMaxLevel down to Layer 0
    for (let level = Math.min(this.maxLevel, nodeMaxLevel); level >= 0; level--) {
      // Find candidate neighbors utilizing greedy BFS search depth (efConstruction)
      const candidates = this.searchLayer(vector, currNode, this.efConstruction, level);
      
      // Connect new node to candidates and enforce max connection limit (M)
      for (const candidate of candidates) {
        newNode.addConnection(level, candidate);
        candidate.addConnection(level, newNode);
        
        // Trim connections if they exceed limit M
        this.shrinkConnections(candidate, level);
      }
      
      currNode = candidates[0] || currNode;
    }

    // Update global entry node if new node has higher level allocation
    if (nodeMaxLevel > this.maxLevel) {
      this.maxLevel = nodeMaxLevel;
      this.enterNode = newNode;
    }
  }

  private searchLayer(
    query: number[], 
    enterPoint: VectorNode, 
    ef: number, 
    level: number
  ): VectorNode[] {
    const visited = new Set<string>([enterPoint.id]);
    const candidates: VectorNode[] = [enterPoint];
    const results: { node: VectorNode; dist: number }[] = [
      { node: enterPoint, dist: this.computeCosine(query, enterPoint.vector) }
    ];

    while (candidates.length > 0) {
      // Sort to evaluate closest candidate node
      candidates.sort((a, b) => this.computeCosine(query, b.vector) - this.computeCosine(query, a.vector));
      const curr = candidates.shift()!;
      
      const neighbors = curr.connections.get(level) || new Set();
      for (const neighbor of neighbors) {
        if (!visited.has(neighbor.id)) {
          visited.add(neighbor.id);
          const dist = this.computeCosine(query, neighbor.vector);
          
          if (dist > results[results.length - 1].dist || results.length < ef) {
            candidates.push(neighbor);
            results.push({ node: neighbor, dist });
            
            // Keep results sorted and bound to ef search depth
            results.sort((a, b) => b.dist - a.dist);
            if (results.length > ef) {
              results.pop();
            }
          }
        }
      }
    }

    return results.map(r => r.node);
  }

  private shrinkConnections(node: VectorNode, level: number) {
    const connections = node.connections.get(level);
    if (!connections || connections.size <= this.M) return;

    // Rank connections by proximity and retain only top M closest neighbors
    const sorted = Array.from(connections).map(conn => ({
      node: conn,
      dist: this.computeCosine(node.vector, conn.vector)
    })).sort((a, b) => b.dist - a.dist);

    const pruned = new Set<VectorNode>();
    for (let i = 0; i < this.M; i++) {
      pruned.add(sorted[i].node);
    }
    node.connections.set(level, pruned);
  }
}