From a084d43e457d364194e79810856eafc1875edf1d Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Sat, 15 Aug 2026 16:34:49 +0530 Subject: [PATCH] feat(storage): restore HNSW graph from embeddings_hnsw across restarts - Persist a versioned neighbors payload (version, m, efConstruction, levels) so Restore can validate stored graphs and reject foreign or legacy formats - HNSWIndex.Restore loads vectors + link lists from the graph table, requiring exact membership parity with the embeddings table - SearchHNSW now restores the persisted graph on first use (no recomputation after a restart) and falls back to a full Build when the stored graph is absent, stale, or built with different parameters - BuildHNSWIndex remains the force-rebuild path - Remove now-dead HNSWNode type (superseded by the payload) - Tests: restart restore, stale-graph fallback + re-persist, and parameter-mismatch rejection --- storage/hnsw.go | 130 ++++++++++++++++++++++--- storage/hnsw_restore_test.go | 181 +++++++++++++++++++++++++++++++++++ storage/vectors.go | 34 ++++++- 3 files changed, 328 insertions(+), 17 deletions(-) create mode 100644 storage/hnsw_restore_test.go diff --git a/storage/hnsw.go b/storage/hnsw.go index 92155b3..531638d 100644 --- a/storage/hnsw.go +++ b/storage/hnsw.go @@ -3,6 +3,7 @@ package storage import ( "container/heap" "context" + "database/sql" "encoding/json" "math" "math/rand" @@ -29,11 +30,6 @@ type HNSWIndex struct { rng *rand.Rand } -type HNSWNode struct { - ID string `json:"id"` - Level int `json:"level"` -} - // NewHNSWIndex creates a new HNSW index with the given parameters. func NewHNSWIndex(efConstruction, m, efSearch int) *HNSWIndex { return &HNSWIndex{ @@ -103,7 +99,12 @@ func (h *HNSWIndex) Build(ctx context.Context, s *Store, model string) error { if err != nil { return err } + return h.buildFrom(ctx, s, model, embeddings) +} +// buildFrom constructs the graph from an in-memory embedding set and +// persists it. +func (h *HNSWIndex) buildFrom(ctx context.Context, s *Store, model string, embeddings map[string][]float32) error { h.mu.Lock() h.vectors = embeddings h.graph = make(map[string]map[int][]string) @@ -124,6 +125,108 @@ func (h *HNSWIndex) Build(ctx context.Context, s *Store, model string) error { return h.persist(ctx, s, model) } +// neighborsPayload is the JSON shape persisted in embeddings_hnsw.neighbors. +// Parameters are recorded so a Restore can reject graphs built with a +// different configuration (link lists depend on m and efConstruction). +type neighborsPayload struct { + Version int `json:"version"` + M int `json:"m"` + EfConstruction int `json:"efConstruction"` + Levels map[string][]string `json:"levels"` +} + +const neighborsPayloadVersion = 1 + +// Restore reconstructs the index from the neighbors graph persisted by a +// previous Build/Upsert/Remove cycle, without recomputing distances. It +// returns true when a stored graph was loaded. It returns false (and leaves +// the index empty) when the stored graph is absent, structurally +// incompatible (different parameters, old payload shape), or out of sync +// with the embeddings table — callers should then fall back to Build. +func (h *HNSWIndex) Restore(ctx context.Context, s *Store, model string) (bool, error) { + embeddings, err := s.AllEmbeddings(ctx, model) + if err != nil { + return false, err + } + if len(embeddings) == 0 { + return false, nil // nothing indexed; nothing to restore + } + + qctx, cancel := s.withTimeout(ctx) + defer cancel() + rows, err := s.db.QueryContext(qctx, + `SELECT node_id, neighbors FROM embeddings_hnsw WHERE model=?`, model) + if err != nil { + return false, err + } + defer func() { _ = rows.Close() }() + + graph := make(map[string]map[string][]string) + for rows.Next() { + var nodeID string + var neighbors sql.NullString + if err := rows.Scan(&nodeID, &neighbors); err != nil { + return false, err + } + var payload neighborsPayload + var parsed map[string][]string + if !neighbors.Valid || neighbors.String == "" { + return false, nil // legacy or empty payload — must rebuild + } + if err := json.Unmarshal([]byte(neighbors.String), &payload); err != nil { + return false, nil // unrecognized shape — must rebuild + } + if payload.Version != neighborsPayloadVersion || + payload.M != h.m || payload.EfConstruction != h.efConstruction { + return false, nil // incompatible parameters — must rebuild + } + if parsed = payload.Levels; parsed == nil { + return false, nil + } + graph[nodeID] = parsed + } + if err := rows.Err(); err != nil { + return false, err + } + + // Exact membership match: the graph is only reusable when it covers the + // current embedding set for this model. + if len(graph) != len(embeddings) { + return false, nil + } + for nodeID := range graph { + if _, ok := embeddings[nodeID]; !ok { + return false, nil + } + } + + h.mu.Lock() + defer h.mu.Unlock() + h.vectors = embeddings + h.graph = make(map[string]map[int][]string, len(graph)) + for nodeID, rawLevels := range graph { + levels := make(map[int][]string, len(rawLevels)) + for key, links := range rawLevels { + l, err := strconv.Atoi(key) + if err != nil { + return false, nil // corrupted payload — must rebuild + } + cleaned := make([]string, 0, len(links)) + for _, id := range links { + if _, ok := embeddings[id]; ok { + cleaned = append(cleaned, id) + } + } + levels[l] = cleaned + } + h.graph[nodeID] = levels + } + entry, maxLevel := h.topNodeLocked() + h.entryPoint = entry + h.maxLevel = maxLevel + return true, nil +} + // insert adds a single vector to the graph. Caller must hold h.mu. func (h *HNSWIndex) insert(nodeID string, vec []float32) { if len(h.graph) == 0 { @@ -318,15 +421,16 @@ func (h *HNSWIndex) persistNodes(ctx context.Context, s *Store, model string, id continue } vec := h.vectors[nodeID] - byLevel := make(map[string][]HNSWNode, len(levels)) + byLevel := make(map[string][]string, len(levels)) for l, links := range levels { - nodes := make([]HNSWNode, len(links)) - for i, id := range links { - nodes[i] = HNSWNode{ID: id, Level: l} - } - byLevel[strconv.Itoa(l)] = nodes + byLevel[strconv.Itoa(l)] = links } - neighborsJSON, err := json.Marshal(byLevel) + payload, err := json.Marshal(neighborsPayload{ + Version: neighborsPayloadVersion, + M: h.m, + EfConstruction: h.efConstruction, + Levels: byLevel, + }) if err != nil { return err } @@ -335,7 +439,7 @@ func (h *HNSWIndex) persistNodes(ctx context.Context, s *Store, model string, id VALUES (?, ?, ?, ?, ?) ON CONFLICT(node_id) DO UPDATE SET vector=excluded.vector, neighbors=excluded.neighbors, updated_at=excluded.updated_at - `, nodeID, EncodeVector(vec), model, string(neighborsJSON), time.Now().UTC()) + `, nodeID, EncodeVector(vec), model, string(payload), time.Now().UTC()) if err != nil { return err } diff --git a/storage/hnsw_restore_test.go b/storage/hnsw_restore_test.go new file mode 100644 index 0000000..2255462 --- /dev/null +++ b/storage/hnsw_restore_test.go @@ -0,0 +1,181 @@ +package storage + +import ( + "context" + "fmt" + "testing" +) + +// TestHNSWRestoreAcrossRestart verifies a fresh store on the same database +// restores the persisted HNSW graph instead of recomputing it. +func TestHNSWRestoreAcrossRestart(t *testing.T) { + dir := t.TempDir() + dbPath := dir + "/restart.db" + ctx := context.Background() + model := "restart-model" + + s1, err := NewStore(dbPath) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 10; i++ { + id := fmt.Sprintf("r-%d", i) + if err := s1.CreateNode(ctx, &Node{ + ID: id, Type: "convention", Content: id, + ContentHash: "h-" + id, Scope: "project", Project: "test", + }); err != nil { + t.Fatal(err) + } + vec := make([]float32, 8) + vec[0] = 1.0 + vec[1] = float32(i) * 0.1 + if err := s1.SaveEmbedding(ctx, id, model, vec); err != nil { + t.Fatal(err) + } + } + if err := s1.BuildHNSWIndex(ctx, model); err != nil { + t.Fatal(err) + } + wantIDs, _, err := s1.SearchHNSW(ctx, model, dirVec(t, 0.5), 3, 50) + if err != nil { + t.Fatal(err) + } + // Close releases the process lock; the second store simulates a restart. + if err := s1.Close(); err != nil { + t.Fatal(err) + } + + s2, err := NewStore(dbPath) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s2.Close() }() + + // The index cache is empty in a new process; Restore must load the + // persisted graph without a fallback build. + idx := NewHNSWIndex(200, 16, 50) + restored, err := idx.Restore(ctx, s2, model) + if err != nil { + t.Fatalf("Restore: %v", err) + } + if !restored { + t.Fatal("expected persisted graph to be restorable, got false") + } + + // And the full search path must return identical results after restart. + gotIDs, _, err := s2.SearchHNSW(ctx, model, dirVec(t, 0.5), 3, 50) + if err != nil { + t.Fatal(err) + } + if len(gotIDs) != len(wantIDs) { + t.Fatalf("result count changed across restart: %v vs %v", gotIDs, wantIDs) + } + for i := range wantIDs { + if gotIDs[i] != wantIDs[i] { + t.Errorf("result %d: got %s want %s", i, gotIDs[i], wantIDs[i]) + } + } +} + +// TestHNSWRestoreFallbackOnStaleGraph verifies that a persisted graph which +// no longer matches the embeddings table is rejected and a rebuild takes +// over (via the SearchHNSW fallback path). +func TestHNSWRestoreFallbackOnStaleGraph(t *testing.T) { + dir := t.TempDir() + dbPath := dir + "/stale.db" + ctx := context.Background() + model := "stale-model" + + s1, err := NewStore(dbPath) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 10; i++ { + id := fmt.Sprintf("s-%d", i) + if err := s1.CreateNode(ctx, &Node{ + ID: id, Type: "convention", Content: id, + ContentHash: "h-" + id, Scope: "project", Project: "test", + }); err != nil { + t.Fatal(err) + } + vec := make([]float32, 8) + vec[0] = 1.0 + vec[1] = float32(i) * 0.1 + if err := s1.SaveEmbedding(ctx, id, model, vec); err != nil { + t.Fatal(err) + } + } + if err := s1.BuildHNSWIndex(ctx, model); err != nil { + t.Fatal(err) + } + if err := s1.Close(); err != nil { + t.Fatal(err) + } + + s2, err := NewStore(dbPath) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s2.Close() }() + + // Simulate an out-of-band embedding write that bypassed the index API: + // a row in embeddings with no corresponding graph row. + if err := s2.CreateNode(ctx, &Node{ + ID: "s-ghost", Type: "convention", Content: "s-ghost", + ContentHash: "h-s-ghost", Scope: "project", Project: "test", + }); err != nil { + t.Fatal(err) + } + if _, err := s2.DB().ExecContext(ctx, + `INSERT INTO embeddings(node_id, vector, model) VALUES('s-ghost', X'0000803F0000803F0000803F0000803F0000803F0000803F0000803F0000803F', ?)`, model); err != nil { + t.Fatal(err) + } + + idx := NewHNSWIndex(200, 16, 50) + restored, err := idx.Restore(ctx, s2, model) + if err != nil { + t.Fatalf("Restore: %v", err) + } + if restored { + t.Fatal("expected stale graph to be rejected, but Restore returned true") + } + + // SearchHNSW must fall back to a full build and still return results. + ids, _, err := s2.SearchHNSW(ctx, model, dirVec(t, 0.9), 1, 50) + if err != nil { + t.Fatalf("SearchHNSW after stale graph: %v", err) + } + if len(ids) == 0 { + t.Fatal("expected results after rebuild fallback") + } + + // The fallback build must have re-persisted the graph so the next + // process can restore it again. + idx2 := NewHNSWIndex(200, 16, 50) + restored2, err := idx2.Restore(ctx, s2, model) + if err != nil { + t.Fatalf("Restore after rebuild: %v", err) + } + if !restored2 { + t.Fatal("expected rebuilt graph to be restorable") + } +} + +// TestHNSWRestoreRejectsForeignParameters verifies graphs persisted by an +// index with different construction parameters are rejected. +func TestHNSWRestoreRejectsForeignParameters(t *testing.T) { + s, cleanup := setupStore(t) + defer cleanup() + ctx := context.Background() + model := "param-model" + seedHNSW(t, s, model, 6, 0.1) // persists via Build (m=16, ef=200) + + other := NewHNSWIndex(400, 32, 50) // different construction parameters + restored, err := other.Restore(ctx, s, model) + if err != nil { + t.Fatalf("Restore: %v", err) + } + if restored { + t.Fatal("expected parameter-mismatched graph to be rejected") + } +} diff --git a/storage/vectors.go b/storage/vectors.go index d56b167..47e7b4a 100644 --- a/storage/vectors.go +++ b/storage/vectors.go @@ -198,7 +198,11 @@ func allEmbeddingsQ(ctx context.Context, q queryable, model string) (map[string] return result, rows.Err() } -// BuildHNSWIndex builds an HNSW index for the given embedding model. +// BuildHNSWIndex rebuilds the HNSW index for the given embedding model +// from the embeddings table, ignoring any persisted graph. Use it to force a +// fresh index after out-of-band embedding changes; the search path restores +// persisted graphs automatically and only falls back to a build when the +// stored graph is missing or stale. func (s *Store) BuildHNSWIndex(ctx context.Context, model string) error { idx := NewHNSWIndex(200, 16, 50) // efConstruction=200, m=16, efSearch=50 if err := idx.Build(ctx, s, model); err != nil { @@ -210,6 +214,26 @@ func (s *Store) BuildHNSWIndex(ctx context.Context, model string) error { return nil } +// acquireHNSWIndex loads the index for a model on first use: it restores the +// graph persisted in embeddings_hnsw (covering process restarts) and falls +// back to a full Build when nothing usable is stored. +func (s *Store) acquireHNSWIndex(ctx context.Context, model string) error { + idx := NewHNSWIndex(200, 16, 50) + restored, err := idx.Restore(ctx, s, model) + if err != nil { + return err + } + if !restored { + if err := idx.Build(ctx, s, model); err != nil { + return err + } + } + s.hnswMu.Lock() + s.hnswIndexes[model] = idx + s.hnswMu.Unlock() + return nil +} + // invalidateHNSWIndex drops the built index for a model so that the next // SearchHNSW rebuilds it from the embeddings table. Used after transactional // embedding writes, which bypass the incremental Upsert/Remove path. @@ -219,14 +243,16 @@ func (s *Store) invalidateHNSWIndex(model string) { s.hnswMu.Unlock() } -// SearchHNSW searches the HNSW index for the given model, building it -// on demand if it has not been built yet. +// SearchHNSW searches the HNSW index for the given model, loading it on +// first use: a persisted graph is restored from embeddings_hnsw (so a +// restart does not recompute the index), with a full Build fallback when no +// usable graph is stored. func (s *Store) SearchHNSW(ctx context.Context, model string, query []float32, k, ef int) ([]string, []float32, error) { s.hnswMu.Lock() idx := s.hnswIndexes[model] s.hnswMu.Unlock() if idx == nil { - if err := s.BuildHNSWIndex(ctx, model); err != nil { + if err := s.acquireHNSWIndex(ctx, model); err != nil { return nil, nil, err } s.hnswMu.Lock()