diff --git a/graph.go b/graph.go index b350c4a..958705c 100644 --- a/graph.go +++ b/graph.go @@ -76,17 +76,23 @@ func (s searchCandidate[K]) Less(o searchCandidate[K]) bool { return s.dist < o.dist } -// search returns the layer node closest to the target node -// within the same layer. +// search returns the nearest neighbors of target within this layer. +// Implements HNSW Algorithm 5 (SEARCH-LAYER): the result set is bounded +// by efSearch during exploration; only the k closest are returned. func (n *layerNode[K]) search( - // k is the number of candidates in the result set. + // k is the number of nearest neighbors to return. k int, efSearch int, target Vector, distance DistanceFunc, ) []searchCandidate[K] { - // This is a basic greedy algorithm to find the entry point at the given level - // that is closest to the target node. + // The exploration budget must be at least k, otherwise the ef-bounded + // result set cannot hold all requested neighbors and search would + // return fewer than k results. + if efSearch < k { + efSearch = k + } + candidates := heap.Heap[searchCandidate[K]]{} candidates.Init(make([]searchCandidate[K], 0, efSearch)) candidates.Push( @@ -99,53 +105,63 @@ func (n *layerNode[K]) search( result = heap.Heap[searchCandidate[K]]{} visited = make(map[K]bool) ) - result.Init(make([]searchCandidate[K], 0, k)) + result.Init(make([]searchCandidate[K], 0, efSearch)) // Begin with the entry node in the result set. result.Push(candidates.Min()) visited[n.Key] = true for candidates.Len() > 0 { - var ( - current = candidates.Pop().node - improved = false - ) + current := candidates.Pop() + + // Standard HNSW termination: if the closest remaining candidate + // is farther than the worst in the ef-bounded result set, + // no further improvement is possible. + if result.Len() >= efSearch && current.dist > result.Max().dist { + break + } // We iterate the map in a sorted, deterministic fashion for // tests. - neighborKeys := maps.Keys(current.neighbors) + neighborKeys := maps.Keys(current.node.neighbors) slices.Sort(neighborKeys) for _, neighborID := range neighborKeys { - neighbor := current.neighbors[neighborID] + neighbor := current.node.neighbors[neighborID] if visited[neighborID] { continue } visited[neighborID] = true dist := distance(neighbor.Value, target) - improved = improved || dist < result.Min().dist - if result.Len() < k { - result.Push(searchCandidate[K]{node: neighbor, dist: dist}) + sc := searchCandidate[K]{node: neighbor, dist: dist} + if result.Len() < efSearch { + result.Push(sc) } else if dist < result.Max().dist { result.PopLast() - result.Push(searchCandidate[K]{node: neighbor, dist: dist}) - } - - candidates.Push(searchCandidate[K]{node: neighbor, dist: dist}) - // Always store candidates if we haven't reached the limit. - if candidates.Len() > efSearch { - candidates.PopLast() + result.Push(sc) + } else { + continue } + candidates.Push(sc) } + } - // Termination condition: no improvement in distance and at least - // kMin candidates in the result set. - if !improved && result.Len() >= k { - break + // Return only the k closest from the ef-sized result set. + res := result.Slice() + slices.SortFunc(res, func(a, b searchCandidate[K]) int { + if a.dist < b.dist { + return -1 + } + if a.dist > b.dist { + return 1 } + // Deterministic tiebreaking by key. + return cmp.Compare(a.node.Key, b.node.Key) + }) + if len(res) > k { + res = res[:k] } - - return result.Slice() + return res } func (n *layerNode[K]) replenish(m int, dist DistanceFunc) { diff --git a/graph_test.go b/graph_test.go index bcce34f..8211507 100644 --- a/graph_test.go +++ b/graph_test.go @@ -262,6 +262,32 @@ func TestGraph_RemoveAllNodes(t *testing.T) { } } +func TestGraph_SearchFindsCorrectNearest(t *testing.T) { + // With the old search algorithm, termination was too aggressive (stopped + // when no neighbor beat result.Min) and the result set was bounded by k + // instead of efSearch. This caused the search to miss true nearest + // neighbors, especially with small k and large efSearch. + g := newTestGraph[int]() + for i := 0; i < 100; i++ { + g.Add(Node[int]{Key: i, Value: Vector{float32(i)}}) + } + + // Search for k=1 nearest to 50.5. The answer must be 50 or 51. + results := g.Search(Vector{50.5}, 1) + require.Len(t, results, 1) + require.Contains(t, []int{50, 51}, results[0].Key, + "expected nearest neighbor to 50.5, got key=%d", results[0].Key) + + // Search for k=3 nearest to 0.0. Must return 0, 1, 2. + results = g.Search(Vector{0.0}, 3) + require.Len(t, results, 3) + keys := make([]int, 3) + for i, r := range results { + keys[i] = r.Key + } + require.Subset(t, []int{0, 1, 2}, keys) +} + func TestGraph_DeleteReplenishUsesGraphDistance(t *testing.T) { // replenish() previously hardcoded CosineDistance. After deleting a // node from a EuclideanDistance graph, replenish must use the correct