Skip to content

perf(hnsw): reuse visited workspaces #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhenghaoz merged 1 commit into from
Aug 12, 2026
Merged

perf(hnsw): reuse visited workspaces #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions internal/core/group_by_hnsw.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ func expandHNSWGroups(
publicScore func(score float32) float32,
nodeBetter func(left, right hnswScoredNode) bool,
prefetch func(neighbors []int),
visited *hnswVisited,
) ([]GroupResult, error) {
accumulator := newGroupAccumulator(metric, options.TopKPerGroup)
groups := make(map[string]struct{}, min(options.GroupCount, len(initial)))
Expand Down Expand Up @@ -93,12 +94,12 @@ func expandHNSWGroups(
}

frontier := ailego.NewHeap(nodeBetter)
visited := make([]bool, len(keys))
visited.reset(len(keys))
for _, node := range initial {
if node.position < 0 || node.position >= len(keys) || visited[node.position] {
if node.position < 0 || node.position >= len(keys) || visited.seen(node.position) {
continue
}
visited[node.position] = true
visited.mark(node.position)
frontier.Push(node)
}
for frontier.Len() != 0 && len(groups) < options.GroupCount {
Expand All @@ -111,10 +112,10 @@ func expandHNSWGroups(
prefetch(adjacent)
}
for _, neighbor := range adjacent {
if visited[neighbor] {
if visited.seen(neighbor) {
continue
}
visited[neighbor] = true
visited.mark(neighbor)
score, err := scoreAt(neighbor)
if err != nil {
return nil, fmt.Errorf("core: score HNSW group expansion node %d: %w", neighbor, err)
Expand Down
42 changes: 24 additions & 18 deletions internal/core/hnsw.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,6 +355,8 @@ func (i *HNSWIndex) Neighbors(key uint64, level int) ([]uint64, error) {
}

func (i *HNSWIndex) insertBuiltNode(ctx context.Context, position int) error {
visited := acquireHNSWVisited(len(i.keys))
defer releaseHNSWVisited(visited)
level := i.levels[position]
if i.entryPoint < 0 {
i.entryPoint = position
Expand All @@ -364,7 +366,7 @@ func (i *HNSWIndex) insertBuiltNode(ctx context.Context, position int) error {
entry := i.entryPoint
query := i.vectorAt(position)
for currentLevel := i.maxLevel; currentLevel > level; currentLevel-- {
nearest, err := i.searchHNSWLayer(ctx, query, []int{entry}, 1, currentLevel)
nearest, err := i.searchHNSWLayer(ctx, query, []int{entry}, 1, currentLevel, visited)
if err != nil {
return err
}
Expand All @@ -373,7 +375,7 @@ func (i *HNSWIndex) insertBuiltNode(ctx context.Context, position int) error {
}
}
for currentLevel := min(level, i.maxLevel); currentLevel >= 0; currentLevel-- {
candidates, err := i.searchHNSWLayer(ctx, query, []int{entry}, i.options.EFConstruction, currentLevel)
candidates, err := i.searchHNSWLayer(ctx, query, []int{entry}, i.options.EFConstruction, currentLevel, visited)
if err != nil {
return err
}
Expand Down Expand Up @@ -403,7 +405,7 @@ type hnswScoredNode struct {
score float32
}

func (i *HNSWIndex) searchHNSWLayer(ctx context.Context, query []float32, entries []int, ef, level int) ([]hnswScoredNode, error) {
func (i *HNSWIndex) searchHNSWLayer(ctx context.Context, query []float32, entries []int, ef, level int, visited *hnswVisited) ([]hnswScoredNode, error) {
limit := min(ef, len(i.keys))
if limit <= 0 {
return []hnswScoredNode{}, nil
Expand All @@ -412,17 +414,17 @@ func (i *HNSWIndex) searchHNSWLayer(ctx context.Context, query []float32, entrie
worse := func(left, right hnswScoredNode) bool { return hnswNodeBetter(i.options.Metric, right, left) }
candidates := ailego.NewHeap(better)
results := ailego.NewHeap(worse)
visited := make([]bool, len(i.keys))
visited.reset(len(i.keys))
for _, entry := range entries {
if entry < 0 || entry >= len(i.keys) || i.levels[entry] < level || visited[entry] {
if entry < 0 || entry >= len(i.keys) || i.levels[entry] < level || visited.seen(entry) {
continue
}
score, err := i.computeDistance(query, i.vectorAt(entry))
if err != nil {
return nil, err
}
node := hnswScoredNode{position: entry, score: score}
visited[entry] = true
visited.mark(entry)
candidates.Push(node)
results.Push(node)
}
Expand All @@ -436,10 +438,10 @@ func (i *HNSWIndex) searchHNSWLayer(ctx context.Context, query []float32, entrie
break
}
for _, neighbor := range i.neighbors[current.position][level] {
if visited[neighbor] {
if visited.seen(neighbor) {
continue
}
visited[neighbor] = true
visited.mark(neighbor)
score, err := i.computeDistance(query, i.vectorAt(neighbor))
if err != nil {
return nil, err
Expand Down Expand Up @@ -674,8 +676,10 @@ func (i *HNSWIndex) SearchHNSWGroups(
return nil, err
}
entry := i.entryPoint
visited := acquireHNSWVisited(len(i.keys))
defer releaseHNSWVisited(visited)
for level := i.maxLevel; level > 0; level-- {
nearest, err := i.searchHNSWLayer(ctx, query, []int{entry}, 1, level)
nearest, err := i.searchHNSWLayer(ctx, query, []int{entry}, 1, level, visited)
if err != nil {
return nil, fmt.Errorf("core: descend HNSW group-by level %d: %w", level, err)
}
Expand All @@ -689,7 +693,7 @@ func (i *HNSWIndex) SearchHNSWGroups(
},
EF: options.EF, PrefetchOffset: options.PrefetchOffset, PrefetchLines: options.PrefetchLines,
}
initial, err := i.searchHNSWBase(ctx, query, entry, max(options.EF, candidateCount), searchOptions)
initial, err := i.searchHNSWBase(ctx, query, entry, max(options.EF, candidateCount), searchOptions, visited)
if err != nil {
return nil, err
}
Expand All @@ -704,7 +708,7 @@ func (i *HNSWIndex) SearchHNSWGroups(
}
return expandHNSWGroups(
ctx, i.options.Metric, i.keys, i.neighbors, initial, options.GroupByOptions,
scoreAt, func(score float32) float32 { return score }, groupNodeBetter(i.options.Metric, i.keys), prefetch,
scoreAt, func(score float32) float32 { return score }, groupNodeBetter(i.options.Metric, i.keys), prefetch, visited,
)
}

Expand Down Expand Up @@ -751,8 +755,10 @@ func (i *HNSWIndex) searchHNSW(ctx context.Context, query []float32, options HNS
}

entry := i.entryPoint
visited := acquireHNSWVisited(len(i.keys))
defer releaseHNSWVisited(visited)
for level := i.maxLevel; level > 0; level-- {
nearest, err := i.searchHNSWLayer(ctx, query, []int{entry}, 1, level)
nearest, err := i.searchHNSWLayer(ctx, query, []int{entry}, 1, level, visited)
if err != nil {
return nil, fmt.Errorf("core: descend HNSW level %d: %w", level, err)
}
Expand All @@ -761,7 +767,7 @@ func (i *HNSWIndex) searchHNSW(ctx context.Context, query []float32, options HNS
}
}
capacity := max(options.EF, options.TopK)
candidates, err := i.searchHNSWBase(ctx, query, entry, capacity, options)
candidates, err := i.searchHNSWBase(ctx, query, entry, capacity, options, visited)
if err != nil {
return nil, err
}
Expand All @@ -775,19 +781,19 @@ func (i *HNSWIndex) searchHNSW(ctx context.Context, query []float32, options HNS
return results, nil
}

func (i *HNSWIndex) searchHNSWBase(ctx context.Context, query []float32, entry, capacity int, options HNSWSearchOptions) ([]hnswScoredNode, error) {
func (i *HNSWIndex) searchHNSWBase(ctx context.Context, query []float32, entry, capacity int, options HNSWSearchOptions, visited *hnswVisited) ([]hnswScoredNode, error) {
better := func(left, right hnswScoredNode) bool { return hnswNodeBetter(i.options.Metric, left, right) }
worse := func(left, right hnswScoredNode) bool { return i.hnswResultNodeBetter(right, left) }
frontier := ailego.NewHeap(better)
accepted := ailego.NewHeap(worse)
visited := make([]bool, len(i.keys))
visited.reset(len(i.keys))

score, err := i.computeDistance(query, i.vectorAt(entry))
if err != nil {
return nil, fmt.Errorf("core: score HNSW entry point: %w", err)
}
start := hnswScoredNode{position: entry, score: score}
visited[entry] = true
visited.mark(entry)
frontier.Push(start)
if i.acceptHNSWResult(start, options.SearchOptions) {
accepted.Push(start)
Expand All @@ -805,10 +811,10 @@ func (i *HNSWIndex) searchHNSWBase(ctx context.Context, query []float32, entry,
neighbors := i.neighbors[current.position][0]
prefetchDenseHNSWNeighbors(i.vectors, i.dimension, neighbors, options.PrefetchOffset, options.PrefetchLines)
for _, neighbor := range neighbors {
if visited[neighbor] {
if visited.seen(neighbor) {
continue
}
visited[neighbor] = true
visited.mark(neighbor)
score, err := i.computeDistance(query, i.vectorAt(neighbor))
if err != nil {
return nil, fmt.Errorf("core: score HNSW node %d: %w", neighbor, err)
Expand Down
17 changes: 10 additions & 7 deletions internal/core/hnsw_parallel.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,10 +66,12 @@ func buildParallelHNSW(
}

func (g *parallelHNSWGraph) insert(ctx context.Context, position int) error {
visited := acquireHNSWVisited(len(g.levels))
defer releaseHNSWVisited(visited)
level := g.levels[position]
entry, maxLevel := g.entrySnapshot()
for currentLevel := maxLevel; currentLevel > level; currentLevel-- {
nearest, err := g.searchLayer(ctx, position, []int{entry}, 1, currentLevel)
nearest, err := g.searchLayer(ctx, position, []int{entry}, 1, currentLevel, visited)
if err != nil {
return err
}
Expand All @@ -78,7 +80,7 @@ func (g *parallelHNSWGraph) insert(ctx context.Context, position int) error {
}
}
for currentLevel := min(level, maxLevel); currentLevel >= 0; currentLevel-- {
candidates, err := g.searchLayer(ctx, position, []int{entry}, g.options.EFConstruction, currentLevel)
candidates, err := g.searchLayer(ctx, position, []int{entry}, g.options.EFConstruction, currentLevel, visited)
if err != nil {
return err
}
Expand Down Expand Up @@ -119,6 +121,7 @@ func (g *parallelHNSWGraph) searchLayer(
entries []int,
ef int,
level int,
visited *hnswVisited,
) ([]hnswScoredNode, error) {
limit := min(ef, len(g.levels))
if limit <= 0 {
Expand All @@ -128,17 +131,17 @@ func (g *parallelHNSWGraph) searchLayer(
worse := func(left, right hnswScoredNode) bool { return hnswNodeBetter(g.options.Metric, right, left) }
candidates := ailego.NewHeap(better)
results := ailego.NewHeap(worse)
visited := make([]bool, len(g.levels))
visited.reset(len(g.levels))
for _, entry := range entries {
if entry < 0 || entry >= len(g.levels) || g.levels[entry] < level || visited[entry] {
if entry < 0 || entry >= len(g.levels) || g.levels[entry] < level || visited.seen(entry) {
continue
}
score, err := g.score(query, entry)
if err != nil {
return nil, err
}
node := hnswScoredNode{position: entry, score: score}
visited[entry] = true
visited.mark(entry)
candidates.Push(node)
results.Push(node)
}
Expand All @@ -153,10 +156,10 @@ func (g *parallelHNSWGraph) searchLayer(
}
g.nodeLocks[current.position].RLock()
for _, neighbor := range g.neighbors[current.position][level] {
if visited[neighbor] {
if visited.seen(neighbor) {
continue
}
visited[neighbor] = true
visited.mark(neighbor)
score, err := g.score(query, neighbor)
if err != nil {
g.nodeLocks[current.position].RUnlock()
Expand Down
34 changes: 20 additions & 14 deletions internal/core/hnsw_quantized.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,9 +171,11 @@ func (i *ScalarQuantizedHNSWIndex) SearchHNSWGroups(
scoreAt := func(position int) (float32, error) {
return QuantizedDistance(i.vectors.metric, i.vectors.codes[position], queryCode)
}
visited := acquireHNSWVisited(len(i.vectors.keys))
defer releaseHNSWVisited(visited)
entry := i.base.entryPoint
for level := i.base.maxLevel; level > 0; level-- {
nearest, err := i.searchLayer(ctx, []int{entry}, 1, level, scoreAt)
nearest, err := i.searchLayer(ctx, []int{entry}, 1, level, scoreAt, visited)
if err != nil {
return nil, fmt.Errorf("core: descend scalar-quantized HNSW group-by level %d: %w", level, err)
}
Expand All @@ -187,7 +189,7 @@ func (i *ScalarQuantizedHNSWIndex) SearchHNSWGroups(
},
EF: options.EF, PrefetchOffset: options.PrefetchOffset, PrefetchLines: options.PrefetchLines,
}
initial, err := i.searchBase(ctx, entry, max(options.EF, candidateCount), searchOptions, scoreAt)
initial, err := i.searchBase(ctx, entry, max(options.EF, candidateCount), searchOptions, scoreAt, visited)
if err != nil {
return nil, err
}
Expand All @@ -199,7 +201,7 @@ func (i *ScalarQuantizedHNSWIndex) SearchHNSWGroups(
}
return expandHNSWGroups(
ctx, i.vectors.metric, i.vectors.keys, i.base.neighbors, initial, options.GroupByOptions,
scoreAt, func(score float32) float32 { return score }, groupNodeBetter(i.vectors.metric, i.vectors.keys), prefetch,
scoreAt, func(score float32) float32 { return score }, groupNodeBetter(i.vectors.metric, i.vectors.keys), prefetch, visited,
)
}

Expand Down Expand Up @@ -251,9 +253,11 @@ func (i *ScalarQuantizedHNSWIndex) search(
scoreAt := func(position int) (float32, error) {
return QuantizedDistance(i.vectors.metric, i.vectors.codes[position], queryCode)
}
visited := acquireHNSWVisited(len(i.vectors.keys))
defer releaseHNSWVisited(visited)
entry := i.base.entryPoint
for level := i.base.maxLevel; level > 0; level-- {
nearest, err := i.searchLayer(ctx, []int{entry}, 1, level, scoreAt)
nearest, err := i.searchLayer(ctx, []int{entry}, 1, level, scoreAt, visited)
if err != nil {
return nil, fmt.Errorf("core: descend scalar-quantized HNSW level %d: %w", level, err)
}
Expand All @@ -262,7 +266,7 @@ func (i *ScalarQuantizedHNSWIndex) search(
}
}
capacity := max(options.EF, options.TopK)
candidates, err := i.searchBase(ctx, entry, capacity, options, scoreAt)
candidates, err := i.searchBase(ctx, entry, capacity, options, scoreAt, visited)
if err != nil {
return nil, err
}
Expand All @@ -281,6 +285,7 @@ func (i *ScalarQuantizedHNSWIndex) searchLayer(
entries []int,
ef, level int,
scoreAt func(int) (float32, error),
visited *hnswVisited,
) ([]hnswScoredNode, error) {
limit := min(ef, len(i.vectors.keys))
if limit <= 0 {
Expand All @@ -291,17 +296,17 @@ func (i *ScalarQuantizedHNSWIndex) searchLayer(
worse := func(left, right hnswScoredNode) bool { return hnswNodeBetter(metric, right, left) }
candidates := ailego.NewHeap(better)
results := ailego.NewHeap(worse)
visited := make([]bool, len(i.vectors.keys))
visited.reset(len(i.vectors.keys))
for _, entry := range entries {
if entry < 0 || entry >= len(i.vectors.keys) || i.base.levels[entry] < level || visited[entry] {
if entry < 0 || entry >= len(i.vectors.keys) || i.base.levels[entry] < level || visited.seen(entry) {
continue
}
score, err := scoreAt(entry)
if err != nil {
return nil, err
}
node := hnswScoredNode{position: entry, score: score}
visited[entry] = true
visited.mark(entry)
candidates.Push(node)
results.Push(node)
}
Expand All @@ -315,10 +320,10 @@ func (i *ScalarQuantizedHNSWIndex) searchLayer(
break
}
for _, neighbor := range i.base.neighbors[current.position][level] {
if visited[neighbor] {
if visited.seen(neighbor) {
continue
}
visited[neighbor] = true
visited.mark(neighbor)
score, err := scoreAt(neighbor)
if err != nil {
return nil, err
Expand Down Expand Up @@ -352,20 +357,21 @@ func (i *ScalarQuantizedHNSWIndex) searchBase(
entry, capacity int,
options HNSWSearchOptions,
scoreAt func(int) (float32, error),
visited *hnswVisited,
) ([]hnswScoredNode, error) {
metric := i.vectors.metric
better := func(left, right hnswScoredNode) bool { return hnswNodeBetter(metric, left, right) }
worse := func(left, right hnswScoredNode) bool { return i.resultNodeBetter(right, left) }
frontier := ailego.NewHeap(better)
accepted := ailego.NewHeap(worse)
visited := make([]bool, len(i.vectors.keys))
visited.reset(len(i.vectors.keys))

score, err := scoreAt(entry)
if err != nil {
return nil, fmt.Errorf("core: score scalar-quantized HNSW entry point: %w", err)
}
start := hnswScoredNode{position: entry, score: score}
visited[entry] = true
visited.mark(entry)
frontier.Push(start)
if i.acceptResult(start, options.SearchOptions) {
accepted.Push(start)
Expand All @@ -382,10 +388,10 @@ func (i *ScalarQuantizedHNSWIndex) searchBase(
neighbors := i.base.neighbors[current.position][0]
prefetchQuantizedHNSWNeighbors(i.vectors.codes, neighbors, options.PrefetchOffset, options.PrefetchLines)
for _, neighbor := range neighbors {
if visited[neighbor] {
if visited.seen(neighbor) {
continue
}
visited[neighbor] = true
visited.mark(neighbor)
score, err := scoreAt(neighbor)
if err != nil {
return nil, fmt.Errorf("core: score scalar-quantized HNSW node %d: %w", neighbor, err)
Expand Down
Loading
Loading