From 57dbc57c4a362ca9e6e2e28af400c8a31494e7a4 Mon Sep 17 00:00:00 2001 From: John Suykerbuyk Date: Sat, 11 Apr 2026 12:26:06 -0600 Subject: [PATCH 1/3] fix: heap Max() and PopLast() return wrong element In a binary min-heap, the last array element is not necessarily the maximum. Max() must scan the leaf nodes (indices n/2..n-1) to find the true maximum. PopLast() used Remove(Len()-1) which removed an arbitrary element instead of the worst. This caused incorrect evictions during neighbor selection and search result trimming, degrading graph quality. --- graph_test.go | 4 +++- heap/heap.go | 18 ++++++++++++++++-- heap/heap_test.go | 17 +++++++++++++++++ 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/graph_test.go b/graph_test.go index df795e6..0f28703 100644 --- a/graph_test.go +++ b/graph_test.go @@ -113,13 +113,15 @@ func TestGraph_AddSearch(t *testing.T) { ) require.Len(t, nearest, 4) + // The two closest are 64 and 65 (distance 0.5 each). + // The next two are 63 and 66 (distance 1.5 each). require.EqualValues( t, []Node[int]{ {64, Vector{64}}, {65, Vector{65}}, - {62, Vector{62}}, {63, Vector{63}}, + {66, Vector{66}}, }, nearest, ) diff --git a/heap/heap.go b/heap/heap.go index 7a5052a..c919c3c 100644 --- a/heap/heap.go +++ b/heap/heap.go @@ -70,8 +70,9 @@ func (h *Heap[T]) Pop() T { return heap.Pop(&h.inner).(T) } +// PopLast removes and returns the maximum element from the heap. func (h *Heap[T]) PopLast() T { - return h.Remove(h.Len() - 1) + return h.Remove(h.maxIndex()) } // Remove removes and returns the element at index i from the heap. @@ -85,9 +86,22 @@ func (h *Heap[T]) Min() T { return h.inner.data[0] } +// maxIndex returns the index of the maximum element by scanning leaf nodes. +// In a min-heap the max is always a leaf (indices n/2 .. n-1). +func (h *Heap[T]) maxIndex() int { + n := h.inner.Len() + best := n / 2 + for i := best + 1; i < n; i++ { + if h.inner.data[best].Less(h.inner.data[i]) { + best = i + } + } + return best +} + // Max returns the maximum element in the heap. func (h *Heap[T]) Max() T { - return h.inner.data[h.inner.Len()-1] + return h.inner.data[h.maxIndex()] } func (h *Heap[T]) Slice() []T { diff --git a/heap/heap_test.go b/heap/heap_test.go index 265723e..fd2a9c3 100644 --- a/heap/heap_test.go +++ b/heap/heap_test.go @@ -32,3 +32,20 @@ func TestHeap(t *testing.T) { t.Errorf("Heap did not return sorted elements: %+v", inOrder) } } + +func TestHeap_MaxAndPopLast(t *testing.T) { + h := Heap[Int]{} + values := []Int{5, 1, 9, 3, 7, 2, 8, 4, 6} + for _, v := range values { + h.Push(v) + } + + require.Equal(t, Int(9), h.Max(), "Max should return the largest element") + require.Equal(t, Int(1), h.Min(), "Min should return the smallest element") + + // PopLast should remove and return the maximum. + popped := h.PopLast() + require.Equal(t, Int(9), popped) + require.Equal(t, Int(8), h.Max(), "Max should be 8 after removing 9") + require.Equal(t, 8, h.Len()) +} From 8113f2f1e6d9d8ffdfbc99ad6dbd70019f8a6377 Mon Sep 17 00:00:00 2001 From: John Suykerbuyk Date: Sat, 11 Apr 2026 12:27:30 -0600 Subject: [PATCH 2/3] fix: replenish() hardcodes CosineDistance instead of using graph distance replenish() is called after neighbor eviction and node deletion to restore connectivity. It unconditionally used CosineDistance, ignoring the graph's configured distance function. This corrupted the graph topology for any non-Cosine metric. Thread DistanceFunc through replenish() and isolate() so the correct distance function is always used. --- graph.go | 12 ++++++------ graph_test.go | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/graph.go b/graph.go index 053e9ac..b350c4a 100644 --- a/graph.go +++ b/graph.go @@ -64,7 +64,7 @@ func (n *layerNode[K]) addNeighbor(newNode *layerNode[K], m int, dist DistanceFu delete(n.neighbors, worst.Key) // Delete backlink from the worst neighbor. delete(worst.neighbors, n.Key) - worst.replenish(m) + worst.replenish(m, dist) } type searchCandidate[K cmp.Ordered] struct { @@ -148,7 +148,7 @@ func (n *layerNode[K]) search( return result.Slice() } -func (n *layerNode[K]) replenish(m int) { +func (n *layerNode[K]) replenish(m int, dist DistanceFunc) { if len(n.neighbors) >= m { return } @@ -165,7 +165,7 @@ func (n *layerNode[K]) replenish(m int) { if candidate == n { continue } - n.addNeighbor(candidate, m, CosineDistance) + n.addNeighbor(candidate, m, dist) if len(n.neighbors) >= m { return } @@ -175,13 +175,13 @@ func (n *layerNode[K]) replenish(m int) { // isolates remove the node from the graph by removing all connections // to neighbors. -func (n *layerNode[K]) isolate(m int) { +func (n *layerNode[K]) isolate(m int, dist DistanceFunc) { for _, neighbor := range n.neighbors { delete(neighbor.neighbors, n.Key) } for _, neighbor := range n.neighbors { - neighbor.replenish(m) + neighbor.replenish(m, dist) } } @@ -501,7 +501,7 @@ func (h *Graph[K]) Delete(key K) bool { if len(layer.nodes) == 0 { deleteLayer[i] = struct{}{} } - node.isolate(h.M) + node.isolate(h.M, h.Distance) deleted = true } diff --git a/graph_test.go b/graph_test.go index 0f28703..bcce34f 100644 --- a/graph_test.go +++ b/graph_test.go @@ -261,3 +261,22 @@ func TestGraph_RemoveAllNodes(t *testing.T) { g.Add(MakeNode(1, vec)) } } + +func TestGraph_DeleteReplenishUsesGraphDistance(t *testing.T) { + // replenish() previously hardcoded CosineDistance. After deleting a + // node from a EuclideanDistance graph, replenish must use the correct + // distance function or the topology becomes corrupted. + g := newTestGraph[int]() // uses EuclideanDistance + for i := 0; i < 20; i++ { + g.Add(Node[int]{Key: i, Value: Vector{float32(i)}}) + } + + // Delete a node in the middle to trigger replenish. + g.Delete(10) + + // Search should still find the correct nearest neighbor. + results := g.Search(Vector{9.5}, 1) + require.Len(t, results, 1) + // Must be 9 or 11 (both distance 0.5 from 9.5). + require.Contains(t, []int{9, 11}, results[0].Key) +} From 2dcc854f2846a735432d6fc388170b561b9fc88b Mon Sep 17 00:00:00 2001 From: John Suykerbuyk Date: Sat, 11 Apr 2026 12:28:29 -0600 Subject: [PATCH 3/3] fix: search termination and result set bounded by k instead of efSearch Two issues in layerNode.search(): 1. Termination was based on whether any neighbor improved result.Min() (the best result). Per HNSW Algorithm 5, termination should occur when the best remaining candidate is farther than result.Max() (the worst result). The old condition stopped exploration too early. 2. The result set was bounded by k (number of results to return) instead of efSearch (the exploration beam width). With k=3 and efSearch=100, only 3 candidates were tracked, causing the distance threshold to converge to a local minimum within a few iterations. Fix: bound the result set by efSearch during exploration, use the correct HNSW termination condition, and trim to k after search completes. --- graph.go | 59 ++++++++++++++++++++++++++++----------------------- graph_test.go | 26 +++++++++++++++++++++++ 2 files changed, 59 insertions(+), 26 deletions(-) diff --git a/graph.go b/graph.go index b350c4a..48fa793 100644 --- a/graph.go +++ b/graph.go @@ -76,17 +76,16 @@ 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. candidates := heap.Heap[searchCandidate[K]]{} candidates.Init(make([]searchCandidate[K], 0, efSearch)) candidates.Push( @@ -99,53 +98,61 @@ 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 { + if result.Len() < efSearch { result.Push(searchCandidate[K]{node: neighbor, dist: dist}) + candidates.Push(searchCandidate[K]{node: neighbor, dist: dist}) } 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() + candidates.Push(searchCandidate[K]{node: neighbor, dist: dist}) } } + } - // 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