diff --git a/cmd/vectord/main.go b/cmd/vectord/main.go index b6e76e6..a9b287f 100644 --- a/cmd/vectord/main.go +++ b/cmd/vectord/main.go @@ -274,21 +274,18 @@ func (h *handlers) handleAdd(w http.ResponseWriter, r *http.Request) { return } } + // Pre-validation above is exhaustive (id, dim, finite, zero-norm), + // so BatchAdd takes the write-lock once and pushes the whole batch + // into coder/hnsw via one variadic Graph.Add. Saves N-1 lock + // acquisitions per HTTP batch. + batch := make([]vectord.BatchItem, len(req.Items)) for j, it := range req.Items { - if err := idx.Add(it.ID, it.Vector, it.Metadata); err != nil { - // Vector-validation errors (NaN/Inf, zero-norm under - // cosine) only surface here; pre-validation is intentional - // minimal scope (id + dim only). - if errors.Is(err, vectord.ErrDimensionMismatch) || - strings.Contains(err.Error(), "non-finite") || - strings.Contains(err.Error(), "zero-norm") { - http.Error(w, "items["+strconv.Itoa(j)+"]: "+err.Error(), http.StatusBadRequest) - return - } - slog.Error("add", "name", name, "id", it.ID, "err", err) - http.Error(w, "internal", http.StatusInternalServerError) - return - } + batch[j] = vectord.BatchItem{ID: it.ID, Vector: it.Vector, Metadata: it.Metadata} + } + if err := idx.BatchAdd(batch); err != nil { + slog.Error("batch add", "name", name, "err", err) + http.Error(w, "internal", http.StatusInternalServerError) + return } // One save per batch (post-loop), not per item. Per scrum // O-W4-style discipline: HTTP-batch boundary is the natural unit. diff --git a/internal/vectord/batch_bench_test.go b/internal/vectord/batch_bench_test.go new file mode 100644 index 0000000..d42a474 --- /dev/null +++ b/internal/vectord/batch_bench_test.go @@ -0,0 +1,69 @@ +package vectord + +import ( + "fmt" + "math/rand" + "testing" +) + +// BenchmarkSingleAdd vs BenchmarkBatchAdd quantifies the lock-amortization +// win for the HTTP-batch shape. Same N items, same vectors; one path +// takes the lock N times, the other takes it once. Run with: +// go test ./internal/vectord/ -bench=. -benchmem -benchtime=1x +func BenchmarkSingleAdd(b *testing.B) { + for _, n := range []int{16, 128, 1024} { + b.Run(fmt.Sprintf("N=%d", n), func(b *testing.B) { + items := makeBatch(n, 768) + for i := 0; i < b.N; i++ { + idx := mustIndex(b) + for _, it := range items { + if err := idx.Add(it.ID, it.Vector, it.Metadata); err != nil { + b.Fatalf("Add: %v", err) + } + } + } + }) + } +} + +func BenchmarkBatchAdd(b *testing.B) { + for _, n := range []int{16, 128, 1024} { + b.Run(fmt.Sprintf("N=%d", n), func(b *testing.B) { + items := makeBatch(n, 768) + for i := 0; i < b.N; i++ { + idx := mustIndex(b) + if err := idx.BatchAdd(items); err != nil { + b.Fatalf("BatchAdd: %v", err) + } + } + }) + } +} + +func mustIndex(tb testing.TB) *Index { + tb.Helper() + idx, err := NewIndex(IndexParams{ + Name: "bench", + Dimension: 768, + M: DefaultM, + EfSearch: DefaultEfSearch, + Distance: DistanceCosine, + }) + if err != nil { + tb.Fatalf("NewIndex: %v", err) + } + return idx +} + +func makeBatch(n, dim int) []BatchItem { + rng := rand.New(rand.NewSource(int64(n))) + out := make([]BatchItem, n) + for i := range out { + v := make([]float32, dim) + for j := range v { + v[j] = rng.Float32()*2 - 1 + } + out[i] = BatchItem{ID: fmt.Sprintf("k-%06d", i), Vector: v} + } + return out +} diff --git a/internal/vectord/index.go b/internal/vectord/index.go index b7da62c..30aeb34 100644 --- a/internal/vectord/index.go +++ b/internal/vectord/index.go @@ -225,6 +225,66 @@ func validateVector(vec []float32, distance string) error { return nil } +// BatchItem is one entry in a BatchAdd call. Same per-field +// contract as Add: ID + Vector required, Metadata follows +// upsert-style semantics (nil = leave existing alone). +type BatchItem struct { + ID string + Vector []float32 + Metadata json.RawMessage +} + +// BatchAdd inserts a slice of items under a single write-lock, with +// one variadic call into coder/hnsw's Graph.Add. Net win vs. a loop +// of single Add calls: N→1 lock acquisitions per HTTP batch and one +// variadic library call instead of N. +// +// Contract: items MUST be pre-validated by the caller (id non-empty, +// vector dimension matches, vector finite + non-zero-norm under +// cosine). Pre-validation lives in the HTTP handler so per-item +// error messages stay precise; reproducing it here would force +// position-encoded errors on every consumer. +// +// Intra-batch duplicate IDs are undefined behavior — coder/hnsw's +// internal "node not added" length-invariant fires on the second +// occurrence. Callers must de-dup before calling. The HTTP smoke +// uses unique IDs so this isn't an exercised path; documented for +// future callers. +func (i *Index) BatchAdd(items []BatchItem) error { + if len(items) == 0 { + return nil + } + i.mu.Lock() + defer i.mu.Unlock() + + // Pre-pass: drop any existing IDs so coder/hnsw's variadic Add + // never sees a re-add. Same library-quirk handling as single + // Add — Len()==1 needs a full graph reset because Delete of the + // last node leaves layers[0] entryless. + for _, it := range items { + if _, exists := i.g.Lookup(it.ID); exists { + if i.g.Len() == 1 { + i.resetGraphLocked() + } else { + i.g.Delete(it.ID) + } + } + } + + nodes := make([]hnsw.Node[string], len(items)) + for j, it := range items { + nodes[j] = hnsw.MakeNode(it.ID, it.Vector) + } + i.g.Add(nodes...) + + for _, it := range items { + if it.Metadata != nil { + i.meta[it.ID] = it.Metadata + } + } + return nil +} + // Delete removes id from the index. Returns true if present. func (i *Index) Delete(id string) bool { i.mu.Lock()