fix(vector): Fix similarity-based HNSW search for cosine and dot product metrics #9559

matthewmcneely · 2026-01-08T02:09:16Z

Summary

This PR fixes a critical bug in HNSW vector search where cosine similarity and dot product metrics returned incorrect results. The search algorithm was treating all metrics as distance metrics (lower is better), causing similarity metrics (higher is better) to return the worst matches instead of the best.

Problem

The HNSW implementation had two issues with similarity-based metrics:

Search phase: The candidate heap in persistent_hnsw.go::searchPersistentLayer always used a min-heap, which pops the lowest value first. For similarity metrics where higher values are better, this caused the algorithm to explore the worst candidates first and terminate prematurely.
Edge pruning phase: The helper.go::addNeighbors function used a fixed comparison (>) when pruning edges, which is correct for distance metrics but inverted for similarity metrics. This resulted in keeping the worst edges instead of the best.

Root Cause

The original code assumed all metrics behave like distance metrics:

// Always used min-heap (pops lowest first)
candidateHeap := *buildPersistentHeapByInit(elements)

// Edge pruning always used > comparison
compare: func(i, j uint64) bool {
    return ph.distance_betw(..., i, ...) > ph.distance_betw(..., j, ...)
}

For Euclidean distance, lower values = better matches → min-heap is correct.
For Cosine/DotProduct similarity, higher values = better matches → need max-heap.

Solution

1. Added candidateHeap interface with metric-aware heap selection

type candidateHeap[T c.Float] interface {
    Len() int
    Pop() minPersistentHeapElement[T]
    Push(minPersistentHeapElement[T])
    PopLast() minPersistentHeapElement[T]
}

func buildCandidateHeap[T c.Float](array []minPersistentHeapElement[T], isSimilarityMetric bool) candidateHeap[T] {
    if isSimilarityMetric {
        return &maxHeapWrapper[T]{...}  // Pops highest first
    }
    return &minHeapWrapper[T]{...}      // Pops lowest first
}

2. Added isSimilarityMetric flag to SimilarityType

type SimilarityType[T c.Float] struct {
    // ... existing fields
    isSimilarityMetric bool  // true for cosine, dotproduct; false for euclidean
}

3. Fixed edge pruning comparison in addNeighbors

compare: func(i, j uint64) bool {
    distI := ph.distance_betw(ctx, tc, uuid, i, &inVec, &outVec)
    distJ := ph.distance_betw(ctx, tc, uuid, j, &inVec, &outVec)
    return !ph.simType.isBetterScore(distI, distJ)
}

Files Changed

File	Changes
tok/hnsw/heap.go	Added candidateHeap interface, minHeapWrapper, maxHeapWrapper, and buildCandidateHeap factory
tok/hnsw/helper.go	Added isSimilarityMetric field to SimilarityType; fixed edge pruning comparison
tok/hnsw/persistent_hnsw.go	Updated searchPersistentLayer to use metric-aware candidate heap
tok/hnsw/persistent_hnsw_test.go	Added comprehensive unit tests for heap behavior and search correctness

Testing

Added new tests covering:

TestCandidateHeapMinHeap: Verifies min-heap pops in ascending order
TestCandidateHeapMaxHeap: Verifies max-heap pops in descending order
TestCandidateHeapPushPop: Tests Push/Pop operations for both heap types
TestCandidateHeapPopLast: Tests PopLast for both types
TestSimilarityTypeIsSimilarityMetric: Verifies flag is set correctly for each metric
TestSearchReturnsCorrectOrderForAllMetrics: End-to-end test for Euclidean, Cosine, and DotProduct
TestEdgePruningKeepsBestEdges: Verifies edge pruning keeps best edges for each metric

Performance Note

This fix builds on PR #9514 which corrected the early termination condition. Together, these changes ensure HNSW search explores the correct number of candidates and returns properly ordered results.

Users experiencing slower insert/search times compared to v25.1.0 can tune performance by lowering efConstruction and efSearch parameters when creating your vector indexes.

Lower values trade recall for speed. The default values (efConstruction=128, efSearch=64) prioritize recall.

GenAI Notice

Parts of this implementation and all of the testing was generated using Claude Opus 4.5 (thinking).

Checklist

The PR title follows the
Conventional Commits syntax, leading
with fix:, feat:, chore:, ci:, etc.
Code compiles correctly and linting (via trunk) passes locally
Tests added for new functionality, or regression tests for bug fixes added as applicable

Fixes #9558

Benchmarks

Our BEIR SciFact Information Retrieval Benchmarks now show recall rates close to or exceeding acceptable and excellent performance for all metrics.

============================================================================================================================================
NDCG@k Comparison
============================================================================================================================================

NDCG@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6200
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5200
  Dgraph v25.1.0 (euclidean)                           ███████████████░░░░░░░░░░░░░░░ 0.5000
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2767
  Dgraph v25.1.0 (dotproduct)                          ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2867
  Dgraph staged-fix (euclidean)                        ███████████████░░░░░░░░░░░░░░░ 0.5233
  Dgraph staged-fix (cosine)                           ███████████████░░░░░░░░░░░░░░░ 0.5300
  Dgraph staged-fix (dotproduct)                       ███████████████░░░░░░░░░░░░░░░ 0.5167

NDCG@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6700
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6000
  Dgraph v25.1.0 (euclidean)                           ████████████████░░░░░░░░░░░░░░ 0.5588
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3043
  Dgraph v25.1.0 (dotproduct)                          █████████░░░░░░░░░░░░░░░░░░░░░ 0.3164
  Dgraph staged-fix (euclidean)                        █████████████████░░░░░░░░░░░░░ 0.5918
  Dgraph staged-fix (cosine)                           █████████████████░░░░░░░░░░░░░ 0.5957
  Dgraph staged-fix (dotproduct)                       █████████████████░░░░░░░░░░░░░ 0.5830

NDCG@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6900
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6300
  Dgraph v25.1.0 (euclidean)                           █████████████████░░░░░░░░░░░░░ 0.5858
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3197
  Dgraph v25.1.0 (dotproduct)                          █████████░░░░░░░░░░░░░░░░░░░░░ 0.3290
  Dgraph staged-fix (euclidean)                        ██████████████████░░░░░░░░░░░░ 0.6168
  Dgraph staged-fix (cosine)                           ██████████████████░░░░░░░░░░░░ 0.6240
  Dgraph staged-fix (dotproduct)                       ██████████████████░░░░░░░░░░░░ 0.6081

NDCG@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████░░░░░░░░░ 0.7000
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████░░░░░░░░░░░ 0.6500
  Dgraph v25.1.0 (euclidean)                           ██████████████████░░░░░░░░░░░░ 0.6118
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3305
  Dgraph v25.1.0 (dotproduct)                          ██████████░░░░░░░░░░░░░░░░░░░░ 0.3423
  Dgraph staged-fix (euclidean)                        ███████████████████░░░░░░░░░░░ 0.6461
  Dgraph staged-fix (cosine)                           ███████████████████░░░░░░░░░░░ 0.6505
  Dgraph staged-fix (dotproduct)                       ███████████████████░░░░░░░░░░░ 0.6369

NDCG@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████░░░░░░░░░ 0.7200
  BEIR Acceptable (Acceptable baseline (384-dim))      ████████████████████░░░░░░░░░░ 0.6800
  Dgraph v25.1.0 (euclidean)                           ███████████████████░░░░░░░░░░░ 0.6418
  Dgraph v25.1.0 (cosine)                              ██████████░░░░░░░░░░░░░░░░░░░░ 0.3445
  Dgraph v25.1.0 (dotproduct)                          ██████████░░░░░░░░░░░░░░░░░░░░ 0.3555
  Dgraph staged-fix (euclidean)                        ████████████████████░░░░░░░░░░ 0.6794
  Dgraph staged-fix (cosine)                           ████████████████████░░░░░░░░░░ 0.6849
  Dgraph staged-fix (dotproduct)                       ████████████████████░░░░░░░░░░ 0.6707

============================================================================================================================================
MAP@k Comparison
============================================================================================================================================

MAP@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6000
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5000
  Dgraph v25.1.0 (euclidean)                           ██████████████░░░░░░░░░░░░░░░░ 0.4812
  Dgraph v25.1.0 (cosine)                              ███████░░░░░░░░░░░░░░░░░░░░░░░ 0.2586
  Dgraph v25.1.0 (dotproduct)                          ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2747
  Dgraph staged-fix (euclidean)                        ███████████████░░░░░░░░░░░░░░░ 0.5046
  Dgraph staged-fix (cosine)                           ███████████████░░░░░░░░░░░░░░░ 0.5112
  Dgraph staged-fix (dotproduct)                       ██████████████░░░░░░░░░░░░░░░░ 0.4979

MAP@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████░░░░░░░░░░░ 0.6400
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████░░░░░░░░░░░░░ 0.5700
  Dgraph v25.1.0 (euclidean)                           ████████████████░░░░░░░░░░░░░░ 0.5357
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2883
  Dgraph v25.1.0 (dotproduct)                          █████████░░░░░░░░░░░░░░░░░░░░░ 0.3022
  Dgraph staged-fix (euclidean)                        ████████████████░░░░░░░░░░░░░░ 0.5663
  Dgraph staged-fix (cosine)                           █████████████████░░░░░░░░░░░░░ 0.5707
  Dgraph staged-fix (dotproduct)                       ████████████████░░░░░░░░░░░░░░ 0.5579

MAP@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████░░░░░░░░░░░ 0.6600
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████░░░░░░░░░░░░░ 0.5900
  Dgraph v25.1.0 (euclidean)                           ████████████████░░░░░░░░░░░░░░ 0.5544
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2993
  Dgraph v25.1.0 (dotproduct)                          █████████░░░░░░░░░░░░░░░░░░░░░ 0.3113
  Dgraph staged-fix (euclidean)                        █████████████████░░░░░░░░░░░░░ 0.5838
  Dgraph staged-fix (cosine)                           █████████████████░░░░░░░░░░░░░ 0.5902
  Dgraph staged-fix (dotproduct)                       █████████████████░░░░░░░░░░░░░ 0.5755

MAP@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6700
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6000
  Dgraph v25.1.0 (euclidean)                           █████████████████░░░░░░░░░░░░░ 0.5676
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3045
  Dgraph v25.1.0 (dotproduct)                          █████████░░░░░░░░░░░░░░░░░░░░░ 0.3175
  Dgraph staged-fix (euclidean)                        █████████████████░░░░░░░░░░░░░ 0.5987
  Dgraph staged-fix (cosine)                           ██████████████████░░░░░░░░░░░░ 0.6035
  Dgraph staged-fix (dotproduct)                       █████████████████░░░░░░░░░░░░░ 0.5900

MAP@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6800
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6100
  Dgraph v25.1.0 (euclidean)                           █████████████████░░░░░░░░░░░░░ 0.5746
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3074
  Dgraph v25.1.0 (dotproduct)                          █████████░░░░░░░░░░░░░░░░░░░░░ 0.3203
  Dgraph staged-fix (euclidean)                        ██████████████████░░░░░░░░░░░░ 0.6060
  Dgraph staged-fix (cosine)                           ██████████████████░░░░░░░░░░░░ 0.6113
  Dgraph staged-fix (dotproduct)                       █████████████████░░░░░░░░░░░░░ 0.5977

============================================================================================================================================
RECALL@k Comparison
============================================================================================================================================

Recall@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6000
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5000
  Dgraph v25.1.0 (euclidean)                           ██████████████░░░░░░░░░░░░░░░░ 0.4812
  Dgraph v25.1.0 (cosine)                              ███████░░░░░░░░░░░░░░░░░░░░░░░ 0.2586
  Dgraph v25.1.0 (dotproduct)                          ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2747
  Dgraph staged-fix (euclidean)                        ███████████████░░░░░░░░░░░░░░░ 0.5046
  Dgraph staged-fix (cosine)                           ███████████████░░░░░░░░░░░░░░░ 0.5112
  Dgraph staged-fix (dotproduct)                       ██████████████░░░░░░░░░░░░░░░░ 0.4979

Recall@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████░░░░░░░░░ 0.7300
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████░░░░░░░░░░░ 0.6500
  Dgraph v25.1.0 (euclidean)                           █████████████████░░░░░░░░░░░░░ 0.5984
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3248
  Dgraph v25.1.0 (dotproduct)                          ██████████░░░░░░░░░░░░░░░░░░░░ 0.3377
  Dgraph staged-fix (euclidean)                        ███████████████████░░░░░░░░░░░ 0.6384
  Dgraph staged-fix (cosine)                           ███████████████████░░░░░░░░░░░ 0.6401
  Dgraph staged-fix (dotproduct)                       ██████████████████░░░░░░░░░░░░ 0.6284

Recall@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████████░░░░░░░ 0.7900
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████████░░░░░░░░░ 0.7200
  Dgraph v25.1.0 (euclidean)                           ███████████████████░░░░░░░░░░░ 0.6638
  Dgraph v25.1.0 (cosine)                              ██████████░░░░░░░░░░░░░░░░░░░░ 0.3632
  Dgraph v25.1.0 (dotproduct)                          ███████████░░░░░░░░░░░░░░░░░░░ 0.3697
  Dgraph staged-fix (euclidean)                        ████████████████████░░░░░░░░░░ 0.6988
  Dgraph staged-fix (cosine)                           █████████████████████░░░░░░░░░ 0.7088
  Dgraph staged-fix (dotproduct)                       ████████████████████░░░░░░░░░░ 0.6888

Recall@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████████░░░░░ 0.8400
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████████░░░░░░░ 0.7800
  Dgraph v25.1.0 (euclidean)                           ██████████████████████░░░░░░░░ 0.7368
  Dgraph v25.1.0 (cosine)                              ███████████░░░░░░░░░░░░░░░░░░░ 0.3950
  Dgraph v25.1.0 (dotproduct)                          ████████████░░░░░░░░░░░░░░░░░░ 0.4074
  Dgraph staged-fix (euclidean)                        ███████████████████████░░░░░░░ 0.7808
  Dgraph staged-fix (cosine)                           ███████████████████████░░░░░░░ 0.7834
  Dgraph staged-fix (dotproduct)                       ███████████████████████░░░░░░░ 0.7701

Recall@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████████████░░ 0.9500
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████████████░░░ 0.9000
  Dgraph v25.1.0 (euclidean)                           ██████████████████████████░░░░ 0.8717
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.4589
  Dgraph v25.1.0 (dotproduct)                          █████████████░░░░░░░░░░░░░░░░░ 0.4658
  Dgraph staged-fix (euclidean)                        ████████████████████████████░░ 0.9350
  Dgraph staged-fix (cosine)                           ████████████████████████████░░ 0.9417
  Dgraph staged-fix (dotproduct)                       ███████████████████████████░░░ 0.9250

============================================================================================================================================
PRECISION@k Comparison
============================================================================================================================================

Precision@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6200
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5200
  Dgraph v25.1.0 (euclidean)                           ███████████████░░░░░░░░░░░░░░░ 0.5000
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2767
  Dgraph v25.1.0 (dotproduct)                          ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2867
  Dgraph staged-fix (euclidean)                        ███████████████░░░░░░░░░░░░░░░ 0.5233
  Dgraph staged-fix (cosine)                           ███████████████░░░░░░░░░░░░░░░ 0.5300
  Dgraph staged-fix (dotproduct)                       ███████████████░░░░░░░░░░░░░░░ 0.5167

Precision@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████████████████ 0.2700
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████████████░░░░░ 0.2300
  Dgraph v25.1.0 (euclidean)                           ████████████████████████░░░░░░ 0.2178
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.1211
  Dgraph v25.1.0 (dotproduct)                          █████████████░░░░░░░░░░░░░░░░░ 0.1211
  Dgraph staged-fix (euclidean)                        █████████████████████████░░░░░ 0.2311
  Dgraph staged-fix (cosine)                           █████████████████████████░░░░░ 0.2322
  Dgraph staged-fix (dotproduct)                       █████████████████████████░░░░░ 0.2278

Precision@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████████████████ 0.1800
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████████████░░░░ 0.1600
  Dgraph v25.1.0 (euclidean)                           ████████████████████████░░░░░░ 0.1480
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.0827
  Dgraph v25.1.0 (dotproduct)                          █████████████░░░░░░░░░░░░░░░░░ 0.0807
  Dgraph staged-fix (euclidean)                        █████████████████████████░░░░░ 0.1553
  Dgraph staged-fix (cosine)                           ██████████████████████████░░░░ 0.1573
  Dgraph staged-fix (dotproduct)                       █████████████████████████░░░░░ 0.1533

Precision@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████████████████ 0.1000
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████████████░░░░ 0.0900
  Dgraph v25.1.0 (euclidean)                           █████████████████████████░░░░░ 0.0837
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.0453
  Dgraph v25.1.0 (dotproduct)                          █████████████░░░░░░░░░░░░░░░░░ 0.0447
  Dgraph staged-fix (euclidean)                        ██████████████████████████░░░░ 0.0887
  Dgraph staged-fix (cosine)                           ██████████████████████████░░░░ 0.0887
  Dgraph staged-fix (dotproduct)                       ██████████████████████████░░░░ 0.0873

Precision@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████████████░░ 0.0100
  BEIR Acceptable (Acceptable baseline (384-dim))      ████████████████████████████░░ 0.0100
  Dgraph v25.1.0 (euclidean)                           ███████████████████████████░░░ 0.0100
  Dgraph v25.1.0 (cosine)                              ██████████████░░░░░░░░░░░░░░░░ 0.0053
  Dgraph v25.1.0 (dotproduct)                          ██████████████░░░░░░░░░░░░░░░░ 0.0051
  Dgraph staged-fix (euclidean)                        █████████████████████████████░ 0.0106
  Dgraph staged-fix (cosine)                           ██████████████████████████████ 0.0107
  Dgraph staged-fix (dotproduct)                       █████████████████████████████░ 0.0105

Previous impl assumed lower values were better (distance) which was breaking similarity search (cosine, dotp)

matthewmcneely · 2026-01-08T02:33:13Z

@joelamming As someone who's crawled around in this tokenizer, would you be willing to take a look at this PR for correctness/completeness?

joelamming · 2026-01-10T22:38:18Z

@matthewmcneely -- this is a great find. I ran some side-by-side synthetic self-retrieval tests with very tight ef with today's main against your PR branch and the results really pop. Apples-to-apples shows same harness and same query regime (k=10):

Fresh cosine index built on main: ~94.6–95.2% miss rate (ef=2 and ef=4)
Fresh cosine index built on your branch: ~0.4–0.5% miss rate under the same plans (~200x fewer failures)
Fixed binary still can't salvage a main-built index: ~93% miss rate at ef=2 (graph is already poisoned)

Cosine/dotproduct are effectively inverted in both candidate exploration and construction pruning. This looks like a construction correctness bug so cosine/dotprod deployments will need to rebuild/reindex their HNSW to see the benefit

I've also spotted a couple of small possible completeness/correctness bits around the pruning heap usage/tests. Going to validate a minimal revision tomorrow and will come back with a diff + results if it holds up

joelamming · 2026-01-11T12:14:25Z

@matthewmcneely #9563

@matthewmcneely

…9563) **Description** This PR is a small follow-up to @matthewmcneely's `matthewmcneely/fix-similarity-based-hnsw-indexing` that tightens the construction-time pruning path in `tok/hnsw/helper.go:addNeighbors`. - Makes construction pruning explicitly maintain the best-to-worst top‑k neighbour list capped at `efConstruction` (rather than relying on heap/slice side-effects) - Avoids duplicate/self edges during construction updates - Aligns the pruning unit test with the same helper used by production to reduce test/production drift This is intended to be easy to cherrypick into the [PR branch](#9559) ### Why The directionality fix is the key part (and @matthewmcneely's [PR](#9559) nails that), but the existing pruning loop still had a couple of structure-level issues (heap initialisation/slice truncation semantics) that could make construction behaviour hard to reason about -- especially under similarity metrics ### Testing - `go test ./tok/hnsw -count=1` - Synthetic recall harness (25k vectors) shows massively improved recall vs. `main` (per my [comment](#9559 (comment))) and stable miss-rate and expected build/query behaviour vs. #9559 **Checklist** - [x] The PR title follows the [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/#summary) syntax, leading with `fix:`, `feat:`, `chore:`, `ci:`, etc. - [x] Code compiles correctly and linting (via trunk) passes locally - [x] Tests added for new functionality, or regression tests for bug fixes added as applicable

matthewmcneely · 2026-01-12T19:42:27Z

@joelamming Thanks for reviewing and that PR. My benchmarking shows not only equal or improved recall, but also your rework on the heap has significantly increased insert speed. Good work! Here's the output of the cosine metric for

v25
this PR before your merging in your PR9563
this PR with your PR9563
this PR w/your PR9563 with similar_to search with twice the effort as defined in the index (index at ef-64, search at ef-128)

Other metrics are performing on par (just elided here to reduce noise).

A note about search and insert speed for v25.1: the "excellent" performance is masking the fact that early/incorrect termination was allowing both to perform fast -- of course the indexes were crap and search results were usually poor (as evidenced by the recall stats).

Also note that the benchmark suite is available at https://github.com/dgraph-io/dgraph-benchmarks/tree/main/vector/beir

Loaded 6 result(s) from benchmark_results.csv

============================================================================================================================================
NDCG@k Comparison
============================================================================================================================================

NDCG@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6200
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5200
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2767
  Dgraph v25.2.0-candidate (cosine)                    ███████████████░░░░░░░░░░░░░░░ 0.5300
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████░░░░░░░░░░░░░░░ 0.5233
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████░░░░░░░░░░░░░░░ 0.5300

NDCG@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6700
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6000
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3043
  Dgraph v25.2.0-candidate (cosine)                    █████████████████░░░░░░░░░░░░░ 0.5957
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            █████████████████░░░░░░░░░░░░░ 0.5901
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      █████████████████░░░░░░░░░░░░░ 0.5968

NDCG@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6900
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6300
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3197
  Dgraph v25.2.0-candidate (cosine)                    ██████████████████░░░░░░░░░░░░ 0.6240
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ██████████████████░░░░░░░░░░░░ 0.6166
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ██████████████████░░░░░░░░░░░░ 0.6233

NDCG@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████░░░░░░░░░ 0.7000
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████░░░░░░░░░░░ 0.6500
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3305
  Dgraph v25.2.0-candidate (cosine)                    ███████████████████░░░░░░░░░░░ 0.6505
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████████░░░░░░░░░░░ 0.6443
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████████░░░░░░░░░░░ 0.6510

NDCG@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████░░░░░░░░░ 0.7200
  BEIR Acceptable (Acceptable baseline (384-dim))      ████████████████████░░░░░░░░░░ 0.6800
  Dgraph v25.1.0 (cosine)                              ██████████░░░░░░░░░░░░░░░░░░░░ 0.3445
  Dgraph v25.2.0-candidate (cosine)                    ████████████████████░░░░░░░░░░ 0.6849
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ████████████████████░░░░░░░░░░ 0.6788
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ████████████████████░░░░░░░░░░ 0.6860

============================================================================================================================================
MAP@k Comparison
============================================================================================================================================

MAP@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6000
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5000
  Dgraph v25.1.0 (cosine)                              ███████░░░░░░░░░░░░░░░░░░░░░░░ 0.2586
  Dgraph v25.2.0-candidate (cosine)                    ███████████████░░░░░░░░░░░░░░░ 0.5112
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████░░░░░░░░░░░░░░░ 0.5046
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████░░░░░░░░░░░░░░░ 0.5112

MAP@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████░░░░░░░░░░░ 0.6400
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████░░░░░░░░░░░░░ 0.5700
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2883
  Dgraph v25.2.0-candidate (cosine)                    █████████████████░░░░░░░░░░░░░ 0.5707
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ████████████████░░░░░░░░░░░░░░ 0.5652
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      █████████████████░░░░░░░░░░░░░ 0.5718

MAP@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████░░░░░░░░░░░ 0.6600
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████░░░░░░░░░░░░░ 0.5900
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2993
  Dgraph v25.2.0-candidate (cosine)                    █████████████████░░░░░░░░░░░░░ 0.5902
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            █████████████████░░░░░░░░░░░░░ 0.5835
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      █████████████████░░░░░░░░░░░░░ 0.5902

MAP@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6700
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6000
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3045
  Dgraph v25.2.0-candidate (cosine)                    ██████████████████░░░░░░░░░░░░ 0.6035
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            █████████████████░░░░░░░░░░░░░ 0.5976
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ██████████████████░░░░░░░░░░░░ 0.6042

MAP@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████░░░░░░░░░░ 0.6800
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████░░░░░░░░░░░░ 0.6100
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3074
  Dgraph v25.2.0-candidate (cosine)                    ██████████████████░░░░░░░░░░░░ 0.6113
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ██████████████████░░░░░░░░░░░░ 0.6054
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ██████████████████░░░░░░░░░░░░ 0.6121

============================================================================================================================================
RECALL@k Comparison
============================================================================================================================================

Recall@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6000
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5000
  Dgraph v25.1.0 (cosine)                              ███████░░░░░░░░░░░░░░░░░░░░░░░ 0.2586
  Dgraph v25.2.0-candidate (cosine)                    ███████████████░░░░░░░░░░░░░░░ 0.5112
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████░░░░░░░░░░░░░░░ 0.5046
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████░░░░░░░░░░░░░░░ 0.5112

Recall@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████░░░░░░░░░ 0.7300
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████░░░░░░░░░░░ 0.6500
  Dgraph v25.1.0 (cosine)                              █████████░░░░░░░░░░░░░░░░░░░░░ 0.3248
  Dgraph v25.2.0-candidate (cosine)                    ███████████████████░░░░░░░░░░░ 0.6401
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████████░░░░░░░░░░░ 0.6351
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████████░░░░░░░░░░░ 0.6417

Recall@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████████░░░░░░░ 0.7900
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████████░░░░░░░░░ 0.7200
  Dgraph v25.1.0 (cosine)                              ██████████░░░░░░░░░░░░░░░░░░░░ 0.3632
  Dgraph v25.2.0-candidate (cosine)                    █████████████████████░░░░░░░░░ 0.7088
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ████████████████████░░░░░░░░░░ 0.6988
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      █████████████████████░░░░░░░░░ 0.7054

Recall@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) █████████████████████████░░░░░ 0.8400
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████████░░░░░░░ 0.7800
  Dgraph v25.1.0 (cosine)                              ███████████░░░░░░░░░░░░░░░░░░░ 0.3950
  Dgraph v25.2.0-candidate (cosine)                    ███████████████████████░░░░░░░ 0.7834
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████████████░░░░░░░ 0.7768
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████████████░░░░░░░ 0.7834

Recall@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ████████████████████████████░░ 0.9500
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████████████░░░ 0.9000
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.4589
  Dgraph v25.2.0-candidate (cosine)                    ████████████████████████████░░ 0.9417
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ████████████████████████████░░ 0.9350
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ████████████████████████████░░ 0.9450

============================================================================================================================================
PRECISION@k Comparison
============================================================================================================================================

Precision@1:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████░░░░░░░░░░░░ 0.6200
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████░░░░░░░░░░░░░░░ 0.5200
  Dgraph v25.1.0 (cosine)                              ████████░░░░░░░░░░░░░░░░░░░░░░ 0.2767
  Dgraph v25.2.0-candidate (cosine)                    ███████████████░░░░░░░░░░░░░░░ 0.5300
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ███████████████░░░░░░░░░░░░░░░ 0.5233
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ███████████████░░░░░░░░░░░░░░░ 0.5300

Precision@3:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████████████████ 0.2700
  BEIR Acceptable (Acceptable baseline (384-dim))      █████████████████████████░░░░░ 0.2300
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.1211
  Dgraph v25.2.0-candidate (cosine)                    █████████████████████████░░░░░ 0.2322
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            █████████████████████████░░░░░ 0.2300
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      █████████████████████████░░░░░ 0.2322

Precision@5:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████████████████ 0.1800
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████████████░░░░ 0.1600
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.0827
  Dgraph v25.2.0-candidate (cosine)                    ██████████████████████████░░░░ 0.1573
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            █████████████████████████░░░░░ 0.1553
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ██████████████████████████░░░░ 0.1567

Precision@10:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ██████████████████████████████ 0.1000
  BEIR Acceptable (Acceptable baseline (384-dim))      ██████████████████████████░░░░ 0.0900
  Dgraph v25.1.0 (cosine)                              █████████████░░░░░░░░░░░░░░░░░ 0.0453
  Dgraph v25.2.0-candidate (cosine)                    ██████████████████████████░░░░ 0.0887
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            ██████████████████████████░░░░ 0.0880
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ██████████████████████████░░░░ 0.0887

Precision@100:
--------------------------------------------------------------------------------------------------------------------------------------------
  BEIR Excellent (State-of-the-art baseline (768-dim)) ███████████████████████████░░░ 0.0100
  BEIR Acceptable (Acceptable baseline (384-dim))      ███████████████████████████░░░ 0.0100
  Dgraph v25.1.0 (cosine)                              ██████████████░░░░░░░░░░░░░░░░ 0.0053
  Dgraph v25.2.0-candidate (cosine)                    █████████████████████████████░ 0.0107
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)            █████████████████████████████░ 0.0106
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef)      ██████████████████████████████ 0.0107

============================================================================================================================================
Timing Comparison
============================================================================================================================================

Insert Time (seconds):
--------------------------------------------------------------------------------------------------------------------------------------------
  Dgraph v25.1.0 (cosine)                         ██████░░░░░░░░░░░░░░░░░░░░░░░░ 20.04s
  Dgraph v25.2.0-candidate (cosine)               ██████████████████████████████ 98.58s
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)       ███████████████████░░░░░░░░░░░ 65.52s
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef) ███████████████████░░░░░░░░░░░ 64.37s

Search Time (seconds):
--------------------------------------------------------------------------------------------------------------------------------------------
  Dgraph v25.1.0 (cosine)                         █████░░░░░░░░░░░░░░░░░░░░░░░░░ 8.21s
  Dgraph v25.2.0-candidate (cosine)               ████████████████████████████░░ 38.61s
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)       ███████████████████████████░░░ 38.45s
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef) ██████████████████████████████ 41.24s

Avg Query Time (milliseconds):
--------------------------------------------------------------------------------------------------------------------------------------------
  Dgraph v25.1.0 (cosine)                         █████░░░░░░░░░░░░░░░░░░░░░░░░░ 27.36ms
  Dgraph v25.2.0-candidate (cosine)               ████████████████████████████░░ 128.68ms
  Dgraph v25.1.0-7-gefe224533 (9563-cosine)       ███████████████████████████░░░ 128.17ms
  Dgraph v25.1.0-7-gefe224533 (9563-cosine-2x-ef) ██████████████████████████████ 137.45ms

minPersistentHeapElement -> persistentHeapElement

Add functionality to correctly evaluate scores

0d519a3

Previous impl assumed lower values were better (distance) which was breaking similarity search (cosine, dotp)

matthewmcneely requested a review from a team as a code owner January 8, 2026 02:09

github-actions bot added area/testing Testing related issues go Pull requests that update Go code labels Jan 8, 2026

matthewmcneely requested review from raphael-istari and shiva-istari January 8, 2026 02:31

Rename min init function for clarity

6fbbcd9

joelamming mentioned this pull request Jan 11, 2026

fix(vector): Make construction pruning deterministic and align tests #9563

Merged

3 tasks

Rename type to better-represent functionality

fb9f0e8

minPersistentHeapElement -> persistentHeapElement

This comment has been minimized.

Sign in to view

mlwelles approved these changes Jan 12, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(vector): Fix similarity-based HNSW search for cosine and dot product metrics #9559

fix(vector): Fix similarity-based HNSW search for cosine and dot product metrics #9559

matthewmcneely commented Jan 8, 2026 •

edited

Loading

Uh oh!

matthewmcneely commented Jan 8, 2026

Uh oh!

joelamming commented Jan 10, 2026

Uh oh!

joelamming commented Jan 11, 2026

Uh oh!

matthewmcneely commented Jan 12, 2026

Uh oh!

This comment has been minimized.

Reviewers

Assignees

Labels

Milestone

Development

Uh oh!

4 participants

fix(vector): Fix similarity-based HNSW search for cosine and dot product metrics #9559

Are you sure you want to change the base?

fix(vector): Fix similarity-based HNSW search for cosine and dot product metrics #9559

Conversation

matthewmcneely commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Problem

Root Cause

Solution

1. Added candidateHeap interface with metric-aware heap selection

2. Added isSimilarityMetric flag to SimilarityType

3. Fixed edge pruning comparison in addNeighbors

Files Changed

Testing

Performance Note

GenAI Notice

Checklist

Benchmarks

Uh oh!

matthewmcneely commented Jan 8, 2026

Uh oh!

joelamming commented Jan 10, 2026

Uh oh!

joelamming commented Jan 11, 2026

Uh oh!

matthewmcneely commented Jan 12, 2026

Uh oh!

This comment has been minimized.

Reviewers

Assignees

Labels

Milestone

Development

Uh oh!

4 participants

matthewmcneely commented Jan 8, 2026 •

edited

Loading