Building a Pure Python HNSW Vector Index with Disk Persistence

# hnsw_index/index.py
import math
import random
import pickle
from typing import List, Tuple, Dict, Set

class HNSWIndex:
    def __init__(self, M: int = 16, ef_construction: int = 200, ef_search: int = 10):
        """
        Initialize the HNSW index.
        :param M: Maximum number of connections per node per layer.
        :param ef_construction: Size of the dynamic candidate list during insertion.
        :param ef_search: Size of the dynamic candidate list during search.
        """
        self.M = M
        self.M_max = M  # Can be tuned for layer 0
        self.ef_construction = ef_construction
        self.ef_search = ef_search
        
        # The graph: a list of dicts, one per vector ID.
        # Each dict has 'vector' and 'neighbors' (a dict of layer -> list of neighbor IDs)
        self.graph: List[Dict] = []
        # Entry point ID (the node with the highest layer)
        self.entry_point: int = -1
        self.max_layer: int = -1
        
    def _sample_level(self) -> int:
        """Sample a layer for a new node using an exponential distribution."""
        # The probability of being assigned to layer l is proportional to 1 / (M^l)
        # We use a simple method: generate a random number and compute the level.
        # Standard HNSW uses: floor(-ln(random()) / ln(M))
        # But we cap at a reasonable max layer to avoid excessive depth.
        max_level = 10  # Practical cap for most datasets
        level = 0
        while random.random() < 1.0 / self.M and level < max_level:
            level += 1
        return level

# Continued in index.py

    def add(self, vector: List[float]) -> int:
        """
        Insert a new vector into the index.
        :param vector: The vector to insert.
        :return: The ID of the inserted vector.
        """
        vec_id = len(self.graph)
        # Store the vector and initialize its neighbors dict
        self.graph.append({
            'vector': vector,
            'neighbors': {}
        })
        
        # Sample a layer for this node
        new_layer = self._sample_level()
        self.graph[vec_id]['neighbors'][new_layer] = []
        
        # If this is the first node, set it as the entry point
        if self.entry_point == -1:
            self.entry_point = vec_id
            self.max_layer = new_layer
            return vec_id
        
        # If the new node's layer is higher than current max, update entry point
        if new_layer > self.max_layer:
            self.entry_point = vec_id
            self.max_layer = new_layer
        
        # Insert the node into the graph layer by layer
        # We start from the top layer and move down to layer 0
        # For each layer, we find the nearest neighbors and connect them.
        current_entry = self.entry_point
        # But we need to traverse from top layer down to new_layer + 1
        # Actually, the standard algorithm: for layer l from max_layer down to new_layer+1, we just search.
        # Then for layers new_layer down to 0, we insert and connect.
        
        # We'll implement a helper for searching in a layer
        # For now, we do a simplified insertion: for each layer from new_layer down to 0,
        # we find the ef_construction nearest neighbors and connect them bidirectionally.
        
        # We need to maintain a candidate set during insertion
        # For simplicity, we do a greedy search at each layer and then connect.
        # This is a simplified version; a full implementation would use a more efficient search.
        
        # Let's do a basic insertion: for each layer l from new_layer down to 0:
        for l in range(new_layer, -1, -1):
            # Find the ef_construction nearest neighbors in this layer
            # We start from the current entry point and traverse the graph at layer l
            neighbors = self._search_layer(vector, current_entry, l, self.ef_construction)
            # Select the M nearest neighbors (or M_max for layer 0)
            m = self.M_max if l == 0 else self.M
            selected = neighbors[:m]
            
            # Connect the new node to selected neighbors
            self.graph[vec_id]['neighbors'][l] = selected
            for neighbor_id in selected:
                if l not in self.graph[neighbor_id]['neighbors']:
                    self.graph[neighbor_id]['neighbors'][l] = []
                self.graph[neighbor_id]['neighbors'][l].append(vec_id)
                # Prune the neighbor's list if it exceeds M
                if len(self.graph[neighbor_id]['neighbors'][l]) > m:
                    # Keep only the M closest ones (we need to compute distances)
                    # For simplicity, we just keep the first M, but this is suboptimal.
                    # In a real implementation, we would sort by distance and prune.
                    self.graph[neighbor_id]['neighbors'][l] = self.graph[neighbor_id]['neighbors'][l][:m]
            
            # Update the entry point for the next layer down
            if selected:
                current_entry = selected[0]  # Use the first neighbor as the new entry for lower layers
        
        return vec_id

    def _search_layer(self, query: List[float], entry_id: int, layer: int, ef: int) -> List[int]:
        """
        Search for the ef nearest neighbors to the query vector in a specific layer.
        Returns a list of node IDs sorted by distance (closest first).
        """
        # We use a set to keep track of visited nodes
        visited: Set[int] = set()
        # We use a list as a priority queue (not efficient, but fine for small ef)
        # We store tuples of (distance, node_id)
        candidates = []
        # Start with the entry point
        entry_dist = self._euclidean_distance(query, self.graph[entry_id]['vector'])
        candidates.append((entry_dist, entry_id))
        visited.add(entry_id)
        
        # Result set: we keep the ef nearest neighbors found
        result = []
        
        while candidates:
            # Pop the closest candidate
            candidates.sort(key=lambda x: x[0])
            current_dist, current_id = candidates.pop(0)
            # Add to result if we haven't reached ef
            if len(result) < ef:
                result.append(current_id)
            else:
                # If the current candidate is farther than the farthest in result, we can stop
                # But we need to check all candidates? Actually, we stop when the closest candidate
                # is farther than the farthest in result.
                # For simplicity, we continue until candidates are empty or we have ef and the next candidate is farther.
                # This is a simplified stopping condition.
                if candidates and candidates[0][0] > current_dist:
                    # Actually, we need to compare with the farthest in result.
                    # Let's compute the max distance in result.
                    max_result_dist = max(self._euclidean_distance(query, self.graph[rid]['vector']) for rid in result)
                    if current_dist > max_result_dist:
                        break
            
            # Explore the neighbors of the current node in this layer
            if layer in self.graph[current_id]['neighbors']:
                for neighbor_id in self.graph[current_id]['neighbors'][layer]:
                    if neighbor_id not in visited:
                        visited.add(neighbor_id)
                        dist = self._euclidean_distance(query, self.graph[neighbor_id]['vector'])
                        candidates.append((dist, neighbor_id))
        
        # Sort result by distance
        result.sort(key=lambda x: self._euclidean_distance(query, self.graph[x]['vector']))
        return result

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

# Continued in index.py

    def search(self, query: List[float], k: int = 1) -> List[Tuple[int, float]]:
        """
        Search for the k nearest neighbors to the query vector.
        Returns a list of (vector_id, distance) tuples.
        """
        if self.entry_point == -1:
            return []
        
        # We start from the top layer and move down to layer 0.
        # At each layer, we search for the ef_search nearest neighbors.
        # We use the closest neighbor from the previous layer as the entry point for the next layer.
        current_entry = self.entry_point
        # We traverse from max_layer down to 0
        for l in range(self.max_layer, -1, -1):
            # Search in this layer with ef = max(ef_search, k)
            ef = max(self.ef_search, k)
            neighbors = self._search_layer(query, current_entry, l, ef)
            # The closest neighbor in this layer becomes the entry for the next layer
            if neighbors:
                current_entry = neighbors[0]
        
        # At layer 0, we do a final search with ef = max(ef_search, k)
        # Actually, we already did layer 0 in the loop, but we can do a more thorough search.
        # For simplicity, we use the result from layer 0.
        # But we need to ensure we have at least k neighbors.
        # Let's do a final search at layer 0 with a larger ef if needed.
        final_neighbors = self._search_layer(query, current_entry, 0, max(self.ef_search, k))
        # Return the top k
        result = [(nid, self._euclidean_distance(query, self.graph[nid]['vector'])) for nid in final_neighbors[:k]]
        return result

    def save(self, filepath: str):
        """Serialize the index to a file using pickle."""
        data = {
            'graph': self.graph,
            'entry_point': self.entry_point,
            'max_layer': self.max_layer,
            'M': self.M,
            'M_max': self.M_max,
            'ef_construction': self.ef_construction,
            'ef_search': self.ef_search
        }
        with open(filepath, 'wb') as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, filepath: str) -> 'HNSWIndex':
        """Load an index from a pickle file."""
        with open(filepath, 'rb') as f:
            data = pickle.load(f)
        index = cls(M=data['M'], ef_construction=data['ef_construction'], ef_search=data['ef_search'])
        index.graph = data['graph']
        index.entry_point = data['entry_point']
        index.max_layer = data['max_layer']
        index.M_max = data['M_max']
        return index

# test_index.py
from hnsw_index.index import HNSWIndex
import random

def test_basic():
    # Create an index
    index = HNSWIndex(M=4, ef_construction=10, ef_search=5)
    
    # Insert some random vectors
    vectors = [[random.random() for _ in range(3)] for _ in range(100)]
    for i, vec in enumerate(vectors):
        index.add(vec)
    
    # Search for the nearest neighbor of the first vector
    query = vectors[0]
    results = index.search(query, k=3)
    print("Search results for vector 0:", results)
    
    # The closest should be itself (distance 0)
    assert results[0][0] == 0, "The query vector should be its own nearest neighbor"
    assert results[0][1] < 1e-6, "Distance to self should be near zero"
    
    # Test persistence
    index.save('test_index.pkl')
    loaded_index = HNSWIndex.load('test_index.pkl')
    
    # Search with the loaded index
    loaded_results = loaded_index.search(query, k=3)
    print("Loaded index search results:", loaded_results)
    
    # Results should be the same
    assert len(loaded_results) == len(results), "Loaded index should return same number of results"
    for (id1, dist1), (id2, dist2) in zip(results, loaded_results):
        assert id1 == id2 and abs(dist1 - dist2) < 1e-6, "Results should match after loading"
    
    print("All tests passed!")

if __name__ == "__main__":
    test_basic()

python test_index.py

import timeit
# Insertion timing
insert_time = timeit.timeit(lambda: index.add([random.random() for _ in range(128)]), number=100)
print(f"Average insertion time: {insert_time/100:.4f} seconds")