fix: search termination and result set bounded by k instead of efSearch

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) {

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