fix: add CountVectors method and optimize ExportVectors logic

Author: krau

internal/infra/imsearch/db/crud.go

@@ -173,6 +173,14 @@ func (d *DB) GetCount(ctx context.Context) (int64, int64, error) {
 	return images, total.Int64, nil
 }
 
+func (d *DB) CountVectors(ctx context.Context) (int64, error) {
+	var n int64
+	err := d.sql.QueryRowContext(ctx,
+		`SELECT COALESCE(SUM(LENGTH(vector)/?), 0) FROM vector`,
+		32).Scan(&n)
+	return n, err
+}
+
 type VectorBound struct {
 	ImageID int64
 	Total   int64 // inclusive upper bound of this image's vector ids

internal/infra/imsearch/imdb/build.go

@@ -11,78 +11,25 @@ import (
 var errNotTrained = errors.New("index backend not trained: run train first")
 
 func (m *IMDB) ExportVectors(ctx context.Context, count int) ([][]byte, error) {
-	imgCount, totalVec, err := m.db.GetCount(ctx)
+	_, totalVec, err := m.db.GetCount(ctx)
 	if err != nil {
 		return nil, err
 	}
-	if imgCount == 0 || totalVec == 0 {
+	if totalVec == 0 {
 		return nil, nil
 	}
 	if int64(count) >= totalVec {
 		return m.exportSequential(ctx, count)
 	}
-
-	avgPerImg := max(int(totalVec/imgCount), 1)
-	imgSample := min(max(count/avgPerImg, 1), int(imgCount))
-
-	out := make([][]byte, 0, count)
-	sqlDB := m.db.SQL()
-
-	rows, err := sqlDB.QueryContext(ctx,
-		`SELECT v.vector FROM vector v
-		 WHERE v.id IN (
-		   SELECT id FROM image
-		   WHERE id IN (SELECT abs(random()) % (SELECT coalesce(max(id),0)+1 FROM image) FROM image LIMIT ?)
-		 )
-		 ORDER BY v.id ASC`, imgSample*3)
-	if err != nil {
-		// Fallback: sequential stride sampling (slow but memory-safe).
-		return m.exportStride(ctx, count, totalVec)
-	}
-	defer rows.Close()
-
-	keepRatio := float64(count) / float64(totalVec)
-	step := 1
-	if keepRatio > 0 && keepRatio < 1 {
-		step = int(1.0 / keepRatio)
+	// totalVec from GetCount is a high-water mark (MAX(total_vector_count)),
+	// not the live vector count after deletes. Use the actual count to compute
+	// the sampling stride so we don't under-sample.
+	actualVec, err := m.db.CountVectors(ctx)
+	if err != nil || actualVec == 0 {
+		actualVec = totalVec
 	}
-	if step < 1 {
-		step = 1
-	}
-
-	for rows.Next() {
-		var blob []byte
-		if err := rows.Scan(&blob); err != nil {
-			return nil, err
-		}
-		n := len(blob) / m.codeSize
-		for i := 0; i < n; i += step {
-			code := make([]byte, m.codeSize)
-			copy(code, blob[i*m.codeSize:(i+1)*m.codeSize])
-			out = append(out, code)
-			if len(out) >= count {
-				return out, nil
-			}
-		}
-	}
-	if len(out) < count {
-		// Random sampling didn't yield enough; fall back to stride.
-		return m.exportStride(ctx, count, totalVec)
-	}
-	return out, rows.Err()
-}
-
-// exportStride reads vectors in id order with a stride, keeping memory bounded.
-func (m *IMDB) exportStride(ctx context.Context, count int, totalVec int64) ([][]byte, error) {
-	keepRatio := float64(count) / float64(totalVec)
-	step := 1
-	if keepRatio > 0 && keepRatio < 1 {
-		step = int(1.0 / keepRatio)
-	}
-	if step < 1 {
-		step = 1
-	}
-	out := make([][]byte, 0, count)
+	keepRatio := float64(count) / float64(actualVec)
+	var out [][]byte
 	const batch = 500
 	var offset int64
 	for len(out) < count {
@@ -95,6 +42,13 @@ func (m *IMDB) exportStride(ctx context.Context, count int, totalVec int64) ([][
 		}
 		for _, r := range rows {
 			n := len(r.Vector) / m.codeSize
+			step := 1
+			if keepRatio > 0 {
+				step = int(1.0 / keepRatio)
+			}
+			if step < 1 {
+				step = 1
+			}
 			for i := 0; i < n; i += step {
 				code := make([]byte, m.codeSize)
 				copy(code, r.Vector[i*m.codeSize:(i+1)*m.codeSize])

internal/infra/imsearch/index/local/local.go

@@ -337,31 +337,15 @@ func (bl *builder) commitShards(nSub int) error {
 		return bl.mergeIntoNew(nSub, finalPath)
 	}
 
-	// invlists.bin exists: merge new shards into it to avoid VStack proliferation.
-	return bl.mergeIntoExisting(nSub, finalPath)
+	// invlists.bin exists: shards are saved alongside it and the searcher
+	// loads them all via VStack. No expensive full-rewrite merge needed.
+	return nil
 }
 
 func (bl *builder) mergeIntoNew(nSub int, finalPath string) error {
-	return bl.mergeShards(nSub, finalPath, false)
-}
-
-func (bl *builder) mergeIntoExisting(nSub int, finalPath string) error {
-	return bl.mergeShards(nSub, finalPath, true)
-}
-
-func (bl *builder) mergeShards(nSub int, finalPath string, includeExisting bool) error {
 	var subs []invlists.InvertedLists
 	var closers []*invlists.OnDisk
 
-	if includeExisting {
-		od, err := invlists.LoadOnDisk(finalPath)
-		if err != nil {
-			return fmt.Errorf("merge: load existing invlists: %w", err)
-		}
-		subs = append(subs, od)
-		closers = append(closers, od)
-	}
-
 	for i := range nSub {
 		od, err := invlists.LoadOnDisk(bl.backend.subIndexPath(i))
 		if err != nil {
@@ -394,7 +378,6 @@ func (bl *builder) mergeShards(nSub int, finalPath string, includeExisting bool)
 	if err := os.Rename(tmpMerged, finalPath); err != nil {
 		return fmt.Errorf("merge: rename: %w", err)
 	}
-	// Remove merged shards.
 	for i := range nSub {
 		os.Remove(bl.backend.subIndexPath(i))
 	}