Score every row, pay for only the ten that survive (V-643)
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YMNNEkYx1mZFtHNrFk7uqb
This commit is contained in:
+135
-11
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user