Merge pull request 'MemoryStore.Search decodes and unmarshals every row before keeping topK' (#195) from task/643-memorystore-search-decodes-and-unmarshal into master
This commit was merged in pull request #195.
This commit is contained in:
@@ -460,8 +460,11 @@ start of a session rather than one lookup per first use:
|
|||||||
ToolSearch("select:mcp__vikunja__list_tasks,mcp__vikunja__get_task_details,mcp__vikunja__create_task,mcp__vikunja__update_task")
|
ToolSearch("select:mcp__vikunja__list_tasks,mcp__vikunja__get_task_details,mcp__vikunja__create_task,mcp__vikunja__update_task")
|
||||||
```
|
```
|
||||||
|
|
||||||
`update_task` carrying a `description` resets `done` to false, so closing a task with a
|
**Close a finished task with `done: true` and nothing else** (owner's call, 07-08-2026).
|
||||||
write-up takes two calls: the description, then `done: true`.
|
Do not write a completion summary into the description on the way out. It is lost anyway,
|
||||||
|
and the durable record is the commit messages and the merged PR. Note that `update_task`
|
||||||
|
carrying a `description` resets `done` to false, which is why a write-up ever took two
|
||||||
|
calls.
|
||||||
|
|
||||||
## Session workflow
|
## Session workflow
|
||||||
|
|
||||||
|
|||||||
+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
|
// Rows under memory.NonRecallPrefix are excluded in SQL. They are speaker
|
||||||
// voiceprints sharing this table, and note recall must not rank them; see that
|
// voiceprints sharing this table, and note recall must not rank them; see that
|
||||||
// constant for why the previous arrangement only appeared to do this.
|
// 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) {
|
func (m *MemoryStore) Search(ctx context.Context, vec []float32, topK int) ([]memory.Result, error) {
|
||||||
if topK <= 0 {
|
if topK <= 0 {
|
||||||
topK = 10
|
topK = 10
|
||||||
@@ -80,30 +87,127 @@ func (m *MemoryStore) Search(ctx context.Context, vec []float32, topK int) ([]me
|
|||||||
}
|
}
|
||||||
defer rows.Close()
|
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() {
|
for rows.Next() {
|
||||||
var id, metaJSON string
|
|
||||||
var blob []byte
|
|
||||||
if err := rows.Scan(&id, &blob, &metaJSON); err != nil {
|
if err := rows.Scan(&id, &blob, &metaJSON); err != nil {
|
||||||
return nil, fmt.Errorf("memory: row: %w", err)
|
return nil, fmt.Errorf("memory: row: %w", err)
|
||||||
}
|
}
|
||||||
meta := map[string]string{}
|
top.offer(dotBlob(vec, blob), id, metaJSON)
|
||||||
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})
|
|
||||||
}
|
}
|
||||||
if err := rows.Err(); err != nil {
|
if err := rows.Err(); err != nil {
|
||||||
return nil, fmt.Errorf("memory: rows: %w", err)
|
return nil, fmt.Errorf("memory: rows: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
sort.Slice(out, func(i, j int) bool { return out[i].Score > out[j].Score })
|
survivors := top.sorted()
|
||||||
if topK < len(out) {
|
out := make([]memory.Result, 0, len(survivors))
|
||||||
out = out[:topK]
|
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
|
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.
|
// ByPrefix returns every row whose id starts with prefix, vectors included.
|
||||||
//
|
//
|
||||||
// This is not a similarity query and deliberately does not score anything:
|
// This is not a similarity query and deliberately does not score anything:
|
||||||
@@ -240,6 +344,26 @@ func decodeVec(b []byte) []float32 {
|
|||||||
return v
|
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 ⇒
|
// dot is the cosine similarity for L2-normalized vectors (mismatched lengths ⇒
|
||||||
// 0, matching internal/memory's cosine).
|
// 0, matching internal/memory's cosine).
|
||||||
func dot(a, b []float32) float64 {
|
func dot(a, b []float32) float64 {
|
||||||
|
|||||||
@@ -0,0 +1,77 @@
|
|||||||
|
package store
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"math"
|
||||||
|
"math/rand"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
// benchDim is the resident embedder's width (multilingual-e5-small, 384), so
|
||||||
|
// the per-row decode cost the benchmark measures is the real one.
|
||||||
|
const benchDim = 384
|
||||||
|
|
||||||
|
// seedMemVectors fills a fresh store with n L2-normalized rows carrying a meta
|
||||||
|
// blob the size recall actually stores — the note text plus its type — because
|
||||||
|
// the cost this benchmark exists to measure is unmarshalling that blob for
|
||||||
|
// every row when only topK survivors need it.
|
||||||
|
func seedMemVectors(tb testing.TB, n int) *MemoryStore {
|
||||||
|
tb.Helper()
|
||||||
|
path := filepath.Join(tb.TempDir(), "mem_bench.db")
|
||||||
|
st, err := Open(context.Background(), path)
|
||||||
|
if err != nil {
|
||||||
|
tb.Fatalf("Open: %v", err)
|
||||||
|
}
|
||||||
|
tb.Cleanup(func() { _ = st.Close() })
|
||||||
|
m := st.VectorMemory()
|
||||||
|
|
||||||
|
rng := rand.New(rand.NewSource(1))
|
||||||
|
ctx := context.Background()
|
||||||
|
for i := 0; i < n; i++ {
|
||||||
|
if err := m.Insert(ctx, fmt.Sprintf("note:%d", i), randUnitVec(rng, benchDim), map[string]string{
|
||||||
|
"type": "note",
|
||||||
|
"text": fmt.Sprintf("заметка номер %d о том, что надо не забыть сделать на неделе", i),
|
||||||
|
}); err != nil {
|
||||||
|
tb.Fatalf("Insert %d: %v", i, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
|
||||||
|
func randUnitVec(rng *rand.Rand, dim int) []float32 {
|
||||||
|
v := make([]float32, dim)
|
||||||
|
var norm float64
|
||||||
|
for i := range v {
|
||||||
|
f := rng.NormFloat64()
|
||||||
|
v[i] = float32(f)
|
||||||
|
norm += f * f
|
||||||
|
}
|
||||||
|
norm = math.Sqrt(norm)
|
||||||
|
for i := range v {
|
||||||
|
v[i] = float32(float64(v[i]) / norm)
|
||||||
|
}
|
||||||
|
return v
|
||||||
|
}
|
||||||
|
|
||||||
|
// BenchmarkMemoryStoreSearch measures one recall query against a store of n
|
||||||
|
// rows. Row counts bracket the documented scale: 1000 is a plausible today,
|
||||||
|
// 10000 is the "thousands, not millions" ceiling the type doc claims a full
|
||||||
|
// scan is fine at.
|
||||||
|
func BenchmarkMemoryStoreSearch(b *testing.B) {
|
||||||
|
for _, n := range []int{1000, 10000} {
|
||||||
|
b.Run(fmt.Sprintf("rows=%d", n), func(b *testing.B) {
|
||||||
|
m := seedMemVectors(b, n)
|
||||||
|
q := randUnitVec(rand.New(rand.NewSource(2)), benchDim)
|
||||||
|
ctx := context.Background()
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if _, err := m.Search(ctx, q, 10); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user