imkk21/custom-vector-search
0
1import numpy as np2import random3import heapq4from typing import Dict, List, Set, Tuple, Optional5from src.metrics import l2_distance, cosine_distance6 7class Node:8 def __init__(self, node_id: str, vector: np.ndarray, level: int):9 self.id = node_id10 self.vector = vector11 self.level = level12 # neighbors[level_idx] = list of neighbor node IDs13 self.neighbors: Dict[int, List[str]] = {l: [] for l in range(level + 1)}14 15class HNSWGraph:16 def __init__(self, M: int = 16, efConstruction: int = 64, efSearch: int = 32, distance_metric: str = "l2"):17 self.M = M18 self.M0 = 2 * M # Max connections at level 019 self.efConstruction = efConstruction20 self.efSearch = efSearch21 self.distance_metric = distance_metric.lower()22 23 # Normalization factor for level generation24 self.mL = 1.0 / np.log(M)25 26 self.nodes: Dict[str, Node] = {}27 self.entry_point: Optional[str] = None28 self.max_level: int = -129 30 def _get_distance(self, v1: np.ndarray, v2: np.ndarray) -> float:31 if self.distance_metric == "cosine":32 return cosine_distance(v1, v2)33 return l2_distance(v1, v2)34 35 def _generate_random_level(self) -> int:36 r = random.random()37 # Avoid log(0)38 if r == 0:39 r = 0.000000140 return int(np.floor(-np.log(r) * self.mL))41 42 def _search_layer(self, query: np.ndarray, enter_points: List[str], ef: int, level: int) -> List[Tuple[float, str]]:43 """44 Finds the ef nearest nodes to query within a single layer, starting from enter_points.45 Returns list of (distance, node_id) tuples.46 """47 # Min-heap of candidates to visit, sorted by distance: (distance, node_id)48 candidates: List[Tuple[float, str]] = []49 # Max-heap of best results found so far: (-distance, node_id)50 # We store negative distance so that heapq's min-heap acts as a max-heap51 v_results: List[Tuple[float, str]] = []52 53 visited: Set[str] = set(enter_points)54 55 for ep in enter_points:56 dist = self._get_distance(query, self.nodes[ep].vector)57 heapq.heappush(candidates, (dist, ep))58 heapq.heappush(v_results, (-dist, ep))59 60 while candidates:61 curr_dist, curr_id = heapq.heappop(candidates)62 63 # If current candidate is further than the furthest result we found, stop searching64 furthest_dist = -v_results[0][0]65 if curr_dist > furthest_dist:66 break67 68 curr_node = self.nodes[curr_id]69 for neighbor_id in curr_node.neighbors.get(level, []):70 if neighbor_id not in visited:71 visited.add(neighbor_id)72 73 neighbor_dist = self._get_distance(query, self.nodes[neighbor_id].vector)74 furthest_dist = -v_results[0][0]75 76 if neighbor_dist < furthest_dist or len(v_results) < ef:77 heapq.heappush(candidates, (neighbor_dist, neighbor_id))78 heapq.heappush(v_results, (-neighbor_dist, neighbor_id))79 80 # Keep results capped at size `ef`81 if len(v_results) > ef:82 heapq.heappop(v_results)83 84 # Convert back to list of (distance, node_id) sorted by distance ascending85 return sorted([( -dist, node_id ) for dist, node_id in v_results], key=lambda x: x[0])86 87 def insert(self, node_id: str, vector: np.ndarray):88 if node_id in self.nodes:89 raise ValueError(f"Node with ID '{node_id}' already exists in graph.")90 91 level = self._generate_random_level()92 new_node = Node(node_id, vector, level)93 self.nodes[node_id] = new_node94 95 if self.entry_point is None:96 self.entry_point = node_id97 self.max_level = level98 return99 100 # Step 1: Find entry point at level of new node by traversing down greedily from max_level101 curr_obj = self.entry_point102 curr_dist = self._get_distance(vector, self.nodes[curr_obj].vector)103 104 # Traverse greedily from max_level down to level + 1105 for l in range(self.max_level, level, -1):106 changed = True107 while changed:108 changed = False109 for neighbor_id in self.nodes[curr_obj].neighbors.get(l, []):110 n_dist = self._get_distance(vector, self.nodes[neighbor_id].vector)111 if n_dist < curr_dist:112 curr_dist = n_dist113 curr_obj = neighbor_id114 changed = True115 116 # Step 2: Insert node at each level from min(level, max_level) down to 0117 enter_points = [curr_obj]118 for l in range(min(level, self.max_level), -1, -1):119 # Find nearest neighbors at this level120 candidates = self._search_layer(vector, enter_points, self.efConstruction, l)121 122 # Select max connections (M at higher levels, M0 at level 0)123 max_conn = self.M0 if l == 0 else self.M124 neighbors_to_connect = candidates[:max_conn]125 126 # Connect the new node to neighbors127 for dist, neighbor_id in neighbors_to_connect:128 new_node.neighbors[l].append(neighbor_id)129 self.nodes[neighbor_id].neighbors[l].append(node_id)130 131 # Shrink connections of neighbor if they exceed maximum limits132 neighbor_max_conn = self.M0 if l == 0 else self.M133 if len(self.nodes[neighbor_id].neighbors[l]) > neighbor_max_conn:134 # Keep the closest ones135 n_vectors = [self.nodes[nid].vector for nid in self.nodes[neighbor_id].neighbors[l]]136 n_dists = [self._get_distance(self.nodes[neighbor_id].vector, n_vec) for n_vec in n_vectors]137 sorted_neighbors = [nid for _, nid in sorted(zip(n_dists, self.nodes[neighbor_id].neighbors[l]))]138 self.nodes[neighbor_id].neighbors[l] = sorted_neighbors[:neighbor_max_conn]139 140 # Update entry points for the next level search141 enter_points = [node_id for _, node_id in candidates]142 143 # Update global entry point if new node level exceeds max_level144 if level > self.max_level:145 self.max_level = level146 self.entry_point = node_id147 148 def search_knn(self, query: np.ndarray, k: int = 5, efSearch: int = None) -> List[Tuple[float, str]]:149 if not self.nodes:150 return []151 152 ef = efSearch if efSearch is not None else self.efSearch153 # Ensure ef is at least k154 ef = max(ef, k)155 156 # Start greedy search from entry point down to level 1157 curr_obj = self.entry_point158 curr_dist = self._get_distance(query, self.nodes[curr_obj].vector)159 160 for l in range(self.max_level, 0, -1):161 changed = True162 while changed:163 changed = False164 for neighbor_id in self.nodes[curr_obj].neighbors.get(l, []):165 n_dist = self._get_distance(query, self.nodes[neighbor_id].vector)166 if n_dist < curr_dist:167 curr_dist = n_dist168 curr_obj = neighbor_id169 changed = True170 171 # At level 0, search layer with efSearch to find closest ef candidates172 results = self._search_layer(query, [curr_obj], ef, 0)173 174 # Return top k results175 return results[:k]176 177 def save(self, filepath: str):178 import pickle179 import os180 os.makedirs(os.path.dirname(filepath), exist_ok=True)181 with open(filepath, "wb") as f:182 pickle.dump(self, f)183 184 @staticmethod185 def load(filepath: str) -> 'HNSWGraph':186 import pickle187 with open(filepath, "rb") as f:188 return pickle.load(f)189