|
| 1 | +//go:build scale |
| 2 | + |
| 3 | +package vectorindex |
| 4 | + |
| 5 | +import ( |
| 6 | + "fmt" |
| 7 | + "math/rand" |
| 8 | + "testing" |
| 9 | +) |
| 10 | + |
| 11 | +// normalizeVec returns v scaled to unit L2 norm (new slice). Lives in |
| 12 | +// this scale-tagged file because TestHNSW_Recall10k is its only caller. |
| 13 | +func normalizeVec(v []float32) []float32 { |
| 14 | + var s float32 |
| 15 | + for _, x := range v { |
| 16 | + s += x * x |
| 17 | + } |
| 18 | + if s == 0 { |
| 19 | + return v |
| 20 | + } |
| 21 | + n := sqrt32(s) |
| 22 | + out := make([]float32, len(v)) |
| 23 | + for i := range v { |
| 24 | + out[i] = v[i] / n |
| 25 | + } |
| 26 | + return out |
| 27 | +} |
| 28 | + |
| 29 | +// TestHNSW_Recall10k builds a 10k-vector index and verifies recall@10 |
| 30 | +// stays above 0.95 across 20 query probes. The workload is fully |
| 31 | +// sequential — the race detector has nothing to catch here — |
| 32 | +// so nightly invokes it WITHOUT -race. Concurrency correctness is |
| 33 | +// covered by TestHNSW_ConcurrentAddSearch, which runs on every PR. |
| 34 | +// |
| 35 | +// Gated behind the `scale` build tag; the nightly workflow runs it via |
| 36 | +// `-tags "sqlite_fts5 scale"`. |
| 37 | +func TestHNSW_Recall10k(t *testing.T) { |
| 38 | + const ( |
| 39 | + n = 10_000 |
| 40 | + dim = 384 |
| 41 | + q = 20 // number of query probes |
| 42 | + k = 10 |
| 43 | + ) |
| 44 | + rng := rand.New(rand.NewSource(7)) |
| 45 | + // Higher construction/search ef for a strong recall benchmark; the |
| 46 | + // default (16/200/50) hits ~0.85 on random vectors which is noisy. |
| 47 | + idx := NewHNSW(32, 400, 400) |
| 48 | + vecs := make(map[string][]float32, n) |
| 49 | + for i := 0; i < n; i++ { |
| 50 | + id := fmt.Sprintf("v%d", i) |
| 51 | + v := normalizeVec(randomVec(rng, dim)) |
| 52 | + vecs[id] = v |
| 53 | + if err := idx.Add(id, v); err != nil { |
| 54 | + t.Fatal(err) |
| 55 | + } |
| 56 | + } |
| 57 | + |
| 58 | + var totalRecall float64 |
| 59 | + for qi := 0; qi < q; qi++ { |
| 60 | + qv := normalizeVec(randomVec(rng, dim)) |
| 61 | + gold := bruteForceTopK(qv, vecs, k) |
| 62 | + hits, err := idx.Search(qv, k) |
| 63 | + if err != nil { |
| 64 | + t.Fatal(err) |
| 65 | + } |
| 66 | + goldSet := map[string]bool{} |
| 67 | + for _, id := range gold { |
| 68 | + goldSet[id] = true |
| 69 | + } |
| 70 | + matches := 0 |
| 71 | + for _, h := range hits { |
| 72 | + if goldSet[h.ID] { |
| 73 | + matches++ |
| 74 | + } |
| 75 | + } |
| 76 | + totalRecall += float64(matches) / float64(k) |
| 77 | + } |
| 78 | + recall := totalRecall / float64(q) |
| 79 | + t.Logf("HNSW recall@10 over %d queries (N=%d, dim=%d) = %.3f", q, n, dim, recall) |
| 80 | + if recall < 0.95 { |
| 81 | + t.Fatalf("recall@10 = %.3f, want >= 0.95", recall) |
| 82 | + } |
| 83 | +} |
0 commit comments