From 65ee0f9c61e10c67e2c12771ea21d132a51fd713 Mon Sep 17 00:00:00 2001 From: claude Date: Fri, 7 Aug 2026 01:48:05 +0400 Subject: [PATCH] Score every row, pay for only the ten that survive (V-643) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Search decoded the vector blob into a []float32 and JSON-unmarshalled the meta map for every row, then sorted all N and threw away everything past topK. Meta only ever matters for a survivor, and the sort answered a question a bounded heap answers cheaper. The scan still visits every row — that is what picks the winners. What it no longer does is allocate for a row it is about to discard. dotBlob reads the vector out of its stored bytes, so scoring costs nothing; a row is copied and its meta unmarshalled only once it has entered the topK. At 10000 rows and topK 10: 70.6ms to 26.8ms, 58MB to 17.5MB, 240k allocs to 60k. Recall is unchanged where it is measured. recall+onnx scores 22/32 with recall@1 70.4% and recall@3 85.2%, identical to before. TestMemoryStoreSearchMatchesNaive pins the ranking against the full-sort implementation it replaced, and TestDotBlobMatchesDot pins bit-identical scores, which the 0.008 gate margin demands. One behaviour did move: ties. sort.Slice is not stable, so equal scores were ordered arbitrarily; the heap now keeps the earliest. Under the real embedder an exact tie is a duplicate vector and nothing moved. Under the hash embedder the eval's floor uses, everything ties at 0 and that run's recall@3 went 74.1% to 81.5% — a number that measures tie order, not retrieval. recall@1 and false recall, the two the eval asserts, are unchanged on both runs. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01YMNNEkYx1mZFtHNrFk7uqb --- internal/store/memory.go | 146 +++++++++++++++++++++++++-- internal/store/memory_search_test.go | 107 ++++++++++++++++++++ 2 files changed, 242 insertions(+), 11 deletions(-) create mode 100644 internal/store/memory_search_test.go diff --git a/internal/store/memory.go b/internal/store/memory.go index e3d886e..3fc6a17 100644 --- a/internal/store/memory.go +++ b/internal/store/memory.go @@ -68,6 +68,13 @@ func (m *MemoryStore) Insert(ctx context.Context, id string, vec []float32, meta // Rows under memory.NonRecallPrefix are excluded in SQL. They are speaker // voiceprints sharing this table, and note recall must not rank them; see that // constant for why the previous arrangement only appeared to do this. +// +// Every row is still scored, because a full scan is what picks the winners. +// What the scan does NOT do is pay for a row it is about to discard: the score +// is read straight off the stored bytes without materializing a []float32, and +// the meta blob is copied and unmarshalled only for a row that has entered the +// topK. Losers cost one dot product and nothing else. Ranking is unchanged — +// same scores, same order, same ties. func (m *MemoryStore) Search(ctx context.Context, vec []float32, topK int) ([]memory.Result, error) { if topK <= 0 { topK = 10 @@ -80,30 +87,127 @@ func (m *MemoryStore) Search(ctx context.Context, vec []float32, topK int) ([]me } defer rows.Close() - var out []memory.Result + // sql.RawBytes hands us the driver's own buffer, valid only until the next + // Next(). Nothing here outlives the row except what topK.offer copies on a + // survivor, so the three columns cost no allocation per row. + var id, blob, metaJSON sql.RawBytes + top := newTopK(topK) for rows.Next() { - var id, metaJSON string - var blob []byte if err := rows.Scan(&id, &blob, &metaJSON); err != nil { return nil, fmt.Errorf("memory: row: %w", err) } - meta := map[string]string{} - if err := json.Unmarshal([]byte(metaJSON), &meta); err != nil { - return nil, fmt.Errorf("memory: unmarshal meta for %q: %w", id, err) - } - out = append(out, memory.Result{ID: id, Score: dot(vec, decodeVec(blob)), Meta: meta}) + top.offer(dotBlob(vec, blob), id, metaJSON) } if err := rows.Err(); err != nil { return nil, fmt.Errorf("memory: rows: %w", err) } - sort.Slice(out, func(i, j int) bool { return out[i].Score > out[j].Score }) - if topK < len(out) { - out = out[:topK] + survivors := top.sorted() + out := make([]memory.Result, 0, len(survivors)) + for _, c := range survivors { + meta := map[string]string{} + if err := json.Unmarshal(c.meta, &meta); err != nil { + return nil, fmt.Errorf("memory: unmarshal meta for %q: %w", c.id, err) + } + out = append(out, memory.Result{ID: c.id, Score: c.score, Meta: meta}) } return out, nil } +// candidate is one row that is currently in the topK: its score, its id, and +// its meta blob copied out of the driver's buffer. The copy is the price of +// surviving, and only survivors pay it. +type candidate struct { + score float64 + id string + meta []byte +} + +// topK keeps the k highest-scoring candidates seen so far as a min-heap, so the +// weakest survivor is always heap[0] and one comparison decides whether a new +// row is worth copying. k is 10 in practice, so the heap is tiny and the whole +// structure fits in cache. +// +// It is a plain slice with hand-written sift operations rather than +// container/heap, because that interface boxes every element into an `any` on +// Push and costs an allocation per surviving row. +type topK struct { + k int + heap []candidate +} + +func newTopK(k int) *topK { + return &topK{k: k, heap: make([]candidate, 0, k)} +} + +// offer admits a row if it beats the weakest survivor, or if the heap is not +// full yet. id and meta are the driver's buffers and are copied here, never +// retained. +// +// A row that only ties the weakest survivor does not displace it, so among +// equal scores the earliest k rows are kept. The full sort this replaced used +// sort.Slice, which is not stable, so it broke such a tie arbitrarily. That is +// the ONE observable difference between the two, and it is deliberate: +// deterministic beats arbitrary. +// +// It is not academic. Under the real embedder an exact tie means duplicate +// vectors and nothing in the recall eval moved (V-643). Under the hash +// embedder the eval's deterministic floor uses, ties are everywhere — it is +// bag-of-words, so every note sharing no word with the query scores exactly 0 +// — and recall@3 on that run moved 74.1% to 81.5% purely because the zeros now +// come out in a fixed order. Neither number measures retrieval. recall@1 and +// false recall, which the eval actually asserts, are unchanged on both runs. +func (t *topK) offer(score float64, id, meta []byte) { + if t.k == 0 { + return + } + if len(t.heap) < t.k { + t.heap = append(t.heap, candidate{score: score, id: string(id), meta: append([]byte(nil), meta...)}) + t.up(len(t.heap) - 1) + return + } + if score <= t.heap[0].score { + return + } + t.heap[0] = candidate{score: score, id: string(id), meta: append([]byte(nil), meta...)} + t.down(0) +} + +func (t *topK) up(i int) { + for i > 0 { + parent := (i - 1) / 2 + if t.heap[parent].score <= t.heap[i].score { + return + } + t.heap[parent], t.heap[i] = t.heap[i], t.heap[parent] + i = parent + } +} + +func (t *topK) down(i int) { + for { + l, r, small := 2*i+1, 2*i+2, i + if l < len(t.heap) && t.heap[l].score < t.heap[small].score { + small = l + } + if r < len(t.heap) && t.heap[r].score < t.heap[small].score { + small = r + } + if small == i { + return + } + t.heap[small], t.heap[i] = t.heap[i], t.heap[small] + i = small + } +} + +// sorted drains the heap into descending score order — what Search returns. +func (t *topK) sorted() []candidate { + out := t.heap + sort.Slice(out, func(i, j int) bool { return out[i].score > out[j].score }) + return out +} + // ByPrefix returns every row whose id starts with prefix, vectors included. // // This is not a similarity query and deliberately does not score anything: @@ -240,6 +344,26 @@ func decodeVec(b []byte) []float32 { return v } +// dotBlob is dot against a vector still in its stored encoding, so scoring a +// row the query is about to discard does not allocate the []float32 that +// decodeVec would build. Same arithmetic, same order of operations, so it +// returns bit-identical scores to dot(a, decodeVec(b)). +// +// A blob whose length isn't a multiple of 4 is truncated to the whole-element +// prefix, matching decodeVec, and a length mismatch is 0, matching dot. +func dotBlob(a []float32, b []byte) float64 { + n := len(b) / 4 + if len(a) != n || n == 0 { + return 0 + } + var sum float64 + for i := 0; i < n; i++ { + f := math.Float32frombits(binary.LittleEndian.Uint32(b[4*i:])) + sum += float64(a[i]) * float64(f) + } + return sum +} + // dot is the cosine similarity for L2-normalized vectors (mismatched lengths ⇒ // 0, matching internal/memory's cosine). func dot(a, b []float32) float64 { diff --git a/internal/store/memory_search_test.go b/internal/store/memory_search_test.go new file mode 100644 index 0000000..2e19a24 --- /dev/null +++ b/internal/store/memory_search_test.go @@ -0,0 +1,107 @@ +package store + +import ( + "context" + "fmt" + "math/rand" + "sort" + "testing" + + "github.com/kami/maven/internal/memory" +) + +// naiveSearch is the implementation Search replaced: score every row into a +// slice, sort the whole slice, truncate. It stays in the test file as the +// reference the bounded-heap version is judged against, because "recall must +// not change" is a claim about output, not about the code that produces it. +func naiveSearch(t *testing.T, m *MemoryStore, vec []float32, topK int) []memory.Result { + t.Helper() + rows, err := m.db.QueryContext(context.Background(), + `SELECT id, vec FROM memory_vectors WHERE id NOT LIKE ? ESCAPE '\'`, + escapeLike(memory.NonRecallPrefix)+"%") + if err != nil { + t.Fatalf("naive scan: %v", err) + } + defer rows.Close() + var out []memory.Result + for rows.Next() { + var id string + var blob []byte + if err := rows.Scan(&id, &blob); err != nil { + t.Fatalf("naive row: %v", err) + } + out = append(out, memory.Result{ID: id, Score: dot(vec, decodeVec(blob))}) + } + if err := rows.Err(); err != nil { + t.Fatalf("naive rows: %v", err) + } + sort.Slice(out, func(i, j int) bool { return out[i].Score > out[j].Score }) + if topK < len(out) { + out = out[:topK] + } + return out +} + +// TestMemoryStoreSearchMatchesNaive is the constraint on V-643: the bounded +// heap must return exactly what a full scan and sort returned. Distinct random +// vectors, so no two scores tie and the ranking is total — a mismatch here is +// arithmetic or heap logic, not a tie-break difference. +func TestMemoryStoreSearchMatchesNaive(t *testing.T) { + ctx := context.Background() + m := newMemTestStore(t).VectorMemory() + + rng := rand.New(rand.NewSource(7)) + const rows, dim = 500, 64 + for i := 0; i < rows; i++ { + if err := m.Insert(ctx, fmt.Sprintf("n%d", i), randUnitVec(rng, dim), map[string]string{ + "text": fmt.Sprintf("note %d", i), + }); err != nil { + t.Fatalf("Insert %d: %v", i, err) + } + } + + for _, topK := range []int{1, 3, 10, 50, rows, rows + 100} { + q := randUnitVec(rng, dim) + got, err := m.Search(ctx, q, topK) + if err != nil { + t.Fatalf("Search topK=%d: %v", topK, err) + } + want := naiveSearch(t, m, q, topK) + if len(got) != len(want) { + t.Fatalf("topK=%d: got %d results, naive returned %d", topK, len(got), len(want)) + } + for i := range want { + if got[i].ID != want[i].ID { + t.Errorf("topK=%d rank %d: got %q, naive says %q", topK, i, got[i].ID, want[i].ID) + } + if got[i].Score != want[i].Score { + t.Errorf("topK=%d rank %d (%s): score %v, naive says %v", + topK, i, got[i].ID, got[i].Score, want[i].Score) + } + } + if len(got) > 0 && got[0].Meta["text"] == "" { + t.Errorf("topK=%d: survivor %s has no meta — it was never unmarshalled", topK, got[0].ID) + } + } +} + +// TestDotBlobMatchesDot pins the claim in dotBlob's doc comment: reading the +// vector out of its stored bytes is bit-identical to decoding it first. Scores +// feed a gate with a 0.008 margin, so "close enough" is not the bar. +func TestDotBlobMatchesDot(t *testing.T) { + rng := rand.New(rand.NewSource(11)) + for i := 0; i < 200; i++ { + a := randUnitVec(rng, 384) + b := randUnitVec(rng, 384) + if got, want := dotBlob(a, encodeVec(b)), dot(a, b); got != want { + t.Fatalf("dotBlob = %v, dot = %v", got, want) + } + } + // Length mismatch is 0 in both, and so is an empty vector. + if got := dotBlob([]float32{1, 0}, encodeVec([]float32{1, 0, 0})); got != 0 { + t.Errorf("mismatched lengths scored %v, want 0", got) + } + if got := dotBlob(nil, nil); got != 0 { + t.Errorf("empty scored %v, want 0", got) + } +}