MB-62182: Avoid re-training vector indexes during merge

centroid_index_test.go

@@ -0,0 +1,74 @@
+//go:build vectors
+// +build vectors
+
+package bleve
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/blevesearch/bleve/v2/analysis/lang/en"
+	"github.com/blevesearch/bleve/v2/mapping"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+func loadSiftData() ([]map[string]interface{}, error) {
+	fileContent, err := os.ReadFile("~/fts/data/datasets/vec-sift-bucket.json")
+	if err != nil {
+		return nil, err
+	}
+	var documents []map[string]interface{}
+	err = json.Unmarshal(fileContent, &documents)
+	if err != nil {
+		return nil, err
+	}
+	return documents, nil
+}
+
+func TestCentroidIndex(t *testing.T) {
+	_, _, err := readDatasetAndQueries(testInputCompressedFile)
+	if err != nil {
+		t.Fatal(err)
+	}
+	documents, err := loadSiftData()
+	if err != nil {
+		t.Fatal(err)
+	}
+	contentFieldMapping := NewTextFieldMapping()
+	contentFieldMapping.Analyzer = en.AnalyzerName
+
+	vecFieldMappingL2 := mapping.NewVectorFieldMapping()
+	vecFieldMappingL2.Dims = 128
+	vecFieldMappingL2.Similarity = index.EuclideanDistance
+
+	indexMappingL2Norm := NewIndexMapping()
+	indexMappingL2Norm.DefaultMapping.AddFieldMappingsAt("content", contentFieldMapping)
+	indexMappingL2Norm.DefaultMapping.AddFieldMappingsAt("vector", vecFieldMappingL2)
+
+	idx, err := newIndexUsing(t.TempDir(), indexMappingL2Norm, Config.DefaultIndexType, Config.DefaultKVStore, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := idx.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	batch := idx.NewBatch()
+	for _, doc := range documents[:100000] {
+		docId := fmt.Sprintf("%s:%s", index.TrainDataPrefix, doc["id"])
+		err = batch.Index(docId, doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	err = idx.Train(batch)
+	if err != nil {
+		t.Fatal(err)
+	}
+}

go.mod

@@ -19,13 +19,13 @@ require (
 	github.com/blevesearch/stempel v0.2.0
 	github.com/blevesearch/upsidedown_store_api v1.0.2
 	github.com/blevesearch/vellum v1.2.0
-	github.com/blevesearch/zapx/v11 v11.4.2
-	github.com/blevesearch/zapx/v12 v12.4.2
-	github.com/blevesearch/zapx/v13 v13.4.2
-	github.com/blevesearch/zapx/v14 v14.4.2
-	github.com/blevesearch/zapx/v15 v15.4.2
-	github.com/blevesearch/zapx/v16 v16.3.0
-	github.com/blevesearch/zapx/v17 v17.0.1
+	github.com/blevesearch/zapx/v11 v11.4.3
+	github.com/blevesearch/zapx/v12 v12.4.3
+	github.com/blevesearch/zapx/v13 v13.4.3
+	github.com/blevesearch/zapx/v14 v14.4.3
+	github.com/blevesearch/zapx/v15 v15.4.3
+	github.com/blevesearch/zapx/v16 v16.3.1
+	github.com/blevesearch/zapx/v17 v17.0.2-0.20260204210735-148661f2ddf6
 	github.com/couchbase/moss v0.2.0
 	github.com/spf13/cobra v1.10.2
 	go.etcd.io/bbolt v1.4.0
@@ -43,3 +43,25 @@ require (
 	github.com/spf13/pflag v1.0.9 // indirect
 	golang.org/x/sys v0.40.0 // indirect
 )
+
+replace github.com/blevesearch/bleve/v2 => /Users/thejas.orkombu/fts/blevesearch/bleve
+
+replace github.com/blevesearch/zapx/v11 => /Users/thejas.orkombu/fts/blevesearch/zapx11
+
+replace github.com/blevesearch/zapx/v12 => /Users/thejas.orkombu/fts/blevesearch/zapx12
+
+replace github.com/blevesearch/zapx/v13 => /Users/thejas.orkombu/fts/blevesearch/zapx13
+
+replace github.com/blevesearch/zapx/v14 => /Users/thejas.orkombu/fts/blevesearch/zapx14
+
+replace github.com/blevesearch/zapx/v15 => /Users/thejas.orkombu/fts/blevesearch/zapx15
+
+replace github.com/blevesearch/zapx/v16 => /Users/thejas.orkombu/fts/blevesearch/zapx
+
+replace github.com/blevesearch/scorch_segment_api/v2 => /Users/thejas.orkombu/fts/blevesearch/scorch_segment_api
+
+replace github.com/blevesearch/go-faiss => /Users/thejas.orkombu/fts/blevesearch/go-faiss
+
+replace github.com/blevesearch/bleve_index_api => /Users/thejas.orkombu/fts/blevesearch/bleve_index_api
+
+replace github.com/blevesearch/sear => /Users/thejas.orkombu/fts/blevesearch/sear

go.sum

@@ -33,20 +33,20 @@ github.com/blevesearch/upsidedown_store_api v1.0.2 h1:U53Q6YoWEARVLd1OYNc9kvhBMG
 github.com/blevesearch/upsidedown_store_api v1.0.2/go.mod h1:M01mh3Gpfy56Ps/UXHjEO/knbqyQ1Oamg8If49gRwrQ=
 github.com/blevesearch/vellum v1.2.0 h1:xkDiOEsHc2t3Cp0NsNZZ36pvc130sCzcGKOPMzXe+e0=
 github.com/blevesearch/vellum v1.2.0/go.mod h1:uEcfBJz7mAOf0Kvq6qoEKQQkLODBF46SINYNkZNae4k=
-github.com/blevesearch/zapx/v11 v11.4.2 h1:l46SV+b0gFN+Rw3wUI1YdMWdSAVhskYuvxlcgpQFljs=
-github.com/blevesearch/zapx/v11 v11.4.2/go.mod h1:4gdeyy9oGa/lLa6D34R9daXNUvfMPZqUYjPwiLmekwc=
-github.com/blevesearch/zapx/v12 v12.4.2 h1:fzRbhllQmEMUuAQ7zBuMvKRlcPA5ESTgWlDEoB9uQNE=
-github.com/blevesearch/zapx/v12 v12.4.2/go.mod h1:TdFmr7afSz1hFh/SIBCCZvcLfzYvievIH6aEISCte58=
-github.com/blevesearch/zapx/v13 v13.4.2 h1:46PIZCO/ZuKZYgxI8Y7lOJqX3Irkc3N8W82QTK3MVks=
-github.com/blevesearch/zapx/v13 v13.4.2/go.mod h1:knK8z2NdQHlb5ot/uj8wuvOq5PhDGjNYQQy0QDnopZk=
-github.com/blevesearch/zapx/v14 v14.4.2 h1:2SGHakVKd+TrtEqpfeq8X+So5PShQ5nW6GNxT7fWYz0=
-github.com/blevesearch/zapx/v14 v14.4.2/go.mod h1:rz0XNb/OZSMjNorufDGSpFpjoFKhXmppH9Hi7a877D8=
-github.com/blevesearch/zapx/v15 v15.4.2 h1:sWxpDE0QQOTjyxYbAVjt3+0ieu8NCE0fDRaFxEsp31k=
-github.com/blevesearch/zapx/v15 v15.4.2/go.mod h1:1pssev/59FsuWcgSnTa0OeEpOzmhtmr/0/11H0Z8+Nw=
-github.com/blevesearch/zapx/v16 v16.3.0 h1:hF6VlN15E9CB40RMPyqOIhlDw1OOo9RItumhKMQktxw=
-github.com/blevesearch/zapx/v16 v16.3.0/go.mod h1:zCFjv7McXWm1C8rROL+3mUoD5WYe2RKsZP3ufqcYpLY=
-github.com/blevesearch/zapx/v17 v17.0.1 h1:kdojyNDiC4abVvsSwequvqYTBuLEXoG3c0UKyxe1+GM=
-github.com/blevesearch/zapx/v17 v17.0.1/go.mod h1:gvr+JMDB9XvQUkT+CaYJhY7aMlez5EmXbkzOBCVyc7U=
+github.com/blevesearch/zapx/v11 v11.4.3 h1:PTZOO5loKpHC/x/GzmPZNa9cw7GZIQxd5qRjwij9tHY=
+github.com/blevesearch/zapx/v11 v11.4.3/go.mod h1:4gdeyy9oGa/lLa6D34R9daXNUvfMPZqUYjPwiLmekwc=
+github.com/blevesearch/zapx/v12 v12.4.3 h1:eElXvAaAX4m04t//CGBQAtHNPA+Q6A1hHZVrN3LSFYo=
+github.com/blevesearch/zapx/v12 v12.4.3/go.mod h1:TdFmr7afSz1hFh/SIBCCZvcLfzYvievIH6aEISCte58=
+github.com/blevesearch/zapx/v13 v13.4.3 h1:qsdhRhaSpVnqDFlRiH9vG5+KJ+dE7KAW9WyZz/KXAiE=
+github.com/blevesearch/zapx/v13 v13.4.3/go.mod h1:knK8z2NdQHlb5ot/uj8wuvOq5PhDGjNYQQy0QDnopZk=
+github.com/blevesearch/zapx/v14 v14.4.3 h1:GY4Hecx0C6UTmiNC2pKdeA2rOKiLR5/rwpU9WR51dgM=
+github.com/blevesearch/zapx/v14 v14.4.3/go.mod h1:rz0XNb/OZSMjNorufDGSpFpjoFKhXmppH9Hi7a877D8=
+github.com/blevesearch/zapx/v15 v15.4.3 h1:iJiMJOHrz216jyO6lS0m9RTCEkprUnzvqAI2lc/0/CU=
+github.com/blevesearch/zapx/v15 v15.4.3/go.mod h1:1pssev/59FsuWcgSnTa0OeEpOzmhtmr/0/11H0Z8+Nw=
+github.com/blevesearch/zapx/v16 v16.3.1 h1:ERxZUSC9UcuKggCQ6b3y4sTkyL4WnGOWuopzglR874g=
+github.com/blevesearch/zapx/v16 v16.3.1/go.mod h1:zCFjv7McXWm1C8rROL+3mUoD5WYe2RKsZP3ufqcYpLY=
+github.com/blevesearch/zapx/v17 v17.0.2-0.20260204210735-148661f2ddf6 h1:eqJh5al0dcPq6VsY6C+G4kva5BBffzMG+sN/SWg2/Eg=
+github.com/blevesearch/zapx/v17 v17.0.2-0.20260204210735-148661f2ddf6/go.mod h1:gvr+JMDB9XvQUkT+CaYJhY7aMlez5EmXbkzOBCVyc7U=
 github.com/couchbase/ghistogram v0.1.0 h1:b95QcQTCzjTUocDXp/uMgSNQi8oj1tGwnJ4bODWZnps=
 github.com/couchbase/ghistogram v0.1.0/go.mod h1:s1Jhy76zqfEecpNWJfWUiKZookAFaiGOEoyzgHt9i7k=
 github.com/couchbase/moss v0.2.0 h1:VCYrMzFwEryyhRSeI+/b3tRBSeTpi/8gn5Kf6dxqn+o=

index.go

@@ -396,3 +396,7 @@ type InsightsIndex interface {
 	// CentroidCardinalities returns the centroids (clusters) from IVF indexes ordered by data density.
 	CentroidCardinalities(field string, limit int, desceding bool) ([]index.CentroidCardinality, error)
 }
+type VectorIndex interface {
+	Index
+	Train(*Batch) error
+}

index/scorch/merge.go

@@ -372,8 +372,8 @@ func (s *Scorch) planMergeAtSnapshot(ctx context.Context,
 
 			atomic.AddUint64(&s.stats.TotFileMergeZapBeg, 1)
 			prevBytesReadTotal := cumulateBytesRead(segmentsToMerge)
-			newDocNums, _, err := s.segPlugin.Merge(segmentsToMerge, docsToDrop, path,
-				cw.cancelCh, s)
+			newDocNums, _, err := s.segPlugin.MergeUsing(segmentsToMerge, docsToDrop, path,
+				cw.cancelCh, s, s.segmentConfig)
 			atomic.AddUint64(&s.stats.TotFileMergeZapEnd, 1)
 
 			fileMergeZapTime := uint64(time.Since(fileMergeZapStartTime))
@@ -391,7 +391,7 @@ func (s *Scorch) planMergeAtSnapshot(ctx context.Context,
 				return fmt.Errorf("merging failed: %v", err)
 			}
 
-			seg, err = s.segPlugin.Open(path)
+			seg, err = s.segPlugin.OpenUsing(path, s.segmentConfig)
 			if err != nil {
 				s.unmarkIneligibleForRemoval(filename)
 				atomic.AddUint64(&s.stats.TotFileMergePlanTasksErr, 1)
@@ -540,7 +540,7 @@ func (s *Scorch) mergeAndPersistInMemorySegments(snapshot *IndexSnapshot,
 			// the newly merged segment is already flushed out to disk, just needs
 			// to be opened using mmap.
 			newDocIDs, _, err :=
-				s.segPlugin.Merge(segsBatch, dropsBatch, path, s.closeCh, s)
+				s.segPlugin.MergeUsing(segsBatch, dropsBatch, path, s.closeCh, s, s.segmentConfig)
 			if err != nil {
 				em.Lock()
 				errs = append(errs, err)
@@ -555,7 +555,7 @@ func (s *Scorch) mergeAndPersistInMemorySegments(snapshot *IndexSnapshot,
 			s.markIneligibleForRemoval(filename)
 			newMergedSegmentIDs[id] = newSegmentID
 			newDocIDsSet[id] = newDocIDs
-			newMergedSegments[id], err = s.segPlugin.Open(path)
+			newMergedSegments[id], err = s.segPlugin.OpenUsing(path, s.segmentConfig)
 			if err != nil {
 				em.Lock()
 				errs = append(errs, err)

index/scorch/persister.go

@@ -425,7 +425,6 @@ func (s *Scorch) persistSnapshotMaybeMerge(snapshot *IndexSnapshot, po *persiste
 	var totSize int
 	var numSegsToFlushOut int
 	var totDocs uint64
-
 	// legacy behaviour of merge + flush of all in-memory segments in one-shot
 	if legacyFlushBehaviour(po.MaxSizeInMemoryMergePerWorker, po.NumPersisterWorkers) {
 		val := &flushable{
@@ -804,7 +803,7 @@ func (s *Scorch) persistSnapshotDirect(snapshot *IndexSnapshot, exclude map[uint
 			}
 		}()
 		for segmentID, path := range newSegmentPaths {
-			newSegments[segmentID], err = s.segPlugin.Open(path)
+			newSegments[segmentID], err = s.segPlugin.OpenUsing(path, s.segmentConfig)
 			if err != nil {
 				return fmt.Errorf("error opening new segment at %s, %v", path, err)
 			}
@@ -853,6 +852,10 @@ func zapFileName(epoch uint64) string {
 	return fmt.Sprintf("%012x.zap", epoch)
 }
 
+func (s *Scorch) loadTrainedData(bucket *bolt.Bucket) error {
+	return s.trainer.loadTrainedData(bucket)
+}
+
 // bolt snapshot code
 
 func (s *Scorch) loadFromBolt() error {
@@ -904,6 +907,12 @@ func (s *Scorch) loadFromBolt() error {
 
 			foundRoot = true
 		}
+
+		trainerBucket := snapshots.Bucket(util.BoltTrainerKey)
+		err := s.trainer.loadTrainedData(trainerBucket)
+		if err != nil {
+			return err
+		}
 		return nil
 	})
 	if err != nil {
@@ -1016,7 +1025,7 @@ func (s *Scorch) loadSegment(segmentBucket *bolt.Bucket) (*SegmentSnapshot, erro
 		return nil, fmt.Errorf("segment path missing")
 	}
 	segmentPath := s.path + string(os.PathSeparator) + string(pathBytes)
-	seg, err := s.segPlugin.Open(segmentPath)
+	seg, err := s.segPlugin.OpenUsing(segmentPath, s.segmentConfig)
 	if err != nil {
 		return nil, fmt.Errorf("error opening bolt segment: %v", err)
 	}

index/scorch/scorch.go

@@ -45,6 +45,7 @@ type Scorch struct {
 	readOnly      bool
 	version       uint8
 	config        map[string]interface{}
+	segmentConfig map[string]interface{}
 	analysisQueue *index.AnalysisQueue
 	path          string
 
@@ -78,6 +79,8 @@ type Scorch struct {
 	rootBolt                 *bolt.DB
 	asyncTasks               sync.WaitGroup
 
+	trainer trainer
+
 	onEvent      func(event Event) bool
 	onAsyncError func(err error, path string)
 
@@ -88,6 +91,29 @@ type Scorch struct {
 	spatialPlugin index.SpatialAnalyzerPlugin
 }
 
+// trainer interface is used for training an index that has the concept
+// of "learning". Naturally, a vector index is one such thing that would
+// implement this interface. There can be multiple implementations of the
+// training itself even for the same index type.
+//
+// this component is not supposed to interact with the other master routines
+// of scorch and will be used only for training the index before the actual data
+// ingestion starts. The routine should also be released once the
+// training is marked as complete - which can be done using the BoltTrainCompleteKey
+// key and a bool value. However the struct is still maintained for the pointer to
+// the instance so that we can use in the later stages of the index lifecycle.
+type trainer interface {
+	// ephemeral
+	trainLoop()
+	// for the training state and the ingestion of the samples
+	train(batch *index.Batch) error
+
+	// to load the metadata from the bolt under the BoltTrainerKey
+	loadTrainedData(*bolt.Bucket) error
+	// to fetch the internal data from the component
+	getInternal(key []byte) ([]byte, error)
+}
+
 type ScorchErrorType string
 
 func (t ScorchErrorType) Error() string {
@@ -154,6 +180,7 @@ func NewScorch(storeName string,
 		forceMergeRequestCh:  make(chan *mergerCtrl, 1),
 		segPlugin:            defaultSegmentPlugin,
 		copyScheduled:        map[string]int{},
+		segmentConfig:        make(map[string]interface{}),
 	}
 
 	forcedSegmentType, forcedSegmentVersion, err := configForceSegmentTypeVersion(config)
@@ -168,6 +195,11 @@ func NewScorch(storeName string,
 		}
 	}
 
+	segConfig, ok := config["segmentConfig"].(map[string]interface{})
+	if ok {
+		rv.segmentConfig = segConfig
+	}
+
 	typ, ok := config["spatialPlugin"].(string)
 	if ok {
 		if err := rv.loadSpatialAnalyzerPlugin(typ); err != nil {
@@ -205,6 +237,8 @@ func NewScorch(storeName string,
 		return nil, err
 	}
 
+	rv.trainer = initTrainer(rv)
+
 	return rv, nil
 }
 
@@ -259,6 +293,9 @@ func (s *Scorch) Open() error {
 	s.asyncTasks.Add(1)
 	go s.introducerLoop()
 
+	s.asyncTasks.Add(1)
+	go s.trainer.trainLoop()
+
 	if !s.readOnly && s.path != "" {
 		s.asyncTasks.Add(1)
 		go s.persisterLoop()
@@ -497,7 +534,7 @@ func (s *Scorch) Batch(batch *index.Batch) (err error) {
 	stats := newFieldStats()
 
 	if len(analysisResults) > 0 {
-		newSegment, bufBytes, err = s.segPlugin.New(analysisResults)
+		newSegment, bufBytes, err = s.segPlugin.NewUsing(analysisResults, s.segmentConfig)
 		if err != nil {
 			return err
 		}
@@ -532,6 +569,21 @@ func (s *Scorch) Batch(batch *index.Batch) (err error) {
 	return err
 }
 
+func (s *Scorch) getInternal(key []byte) ([]byte, error) {
+	s.rootLock.RLock()
+	defer s.rootLock.RUnlock()
+
+	switch string(key) {
+	case string(util.BoltTrainCompleteKey):
+		return s.trainer.getInternal(key)
+	}
+	return nil, nil
+}
+
+func (s *Scorch) Train(batch *index.Batch) error {
+	return s.trainer.train(batch)
+}
+
 func (s *Scorch) prepareSegment(newSegment segment.Segment, ids []string,
 	internalOps map[string][]byte, persistedCallback index.BatchCallback, stats *fieldStats,
 ) error {

index/scorch/segment_plugin.go

@@ -22,12 +22,12 @@
 	index "github.com/blevesearch/bleve_index_api"
 	segment "github.com/blevesearch/scorch_segment_api/v2"
 
 	zapv11 "github.com/blevesearch/zapx/v11"
 	zapv12 "github.com/blevesearch/zapx/v12"
 	zapv13 "github.com/blevesearch/zapx/v13"
 	zapv14 "github.com/blevesearch/zapx/v14"
 	zapv15 "github.com/blevesearch/zapx/v15"
 	zapv16 "github.com/blevesearch/zapx/v16"
 	zapv17 "github.com/blevesearch/zapx/v17"
 )
 
@@ -46,10 +46,14 @@
 	// New takes a set of Documents and turns them into a new Segment
 	New(results []index.Document) (segment.Segment, uint64, error)
 
+	NewUsing(results []index.Document, config map[string]interface{}) (segment.Segment, uint64, error)
+
 	// Open attempts to open the file at the specified path and
 	// return the corresponding Segment
 	Open(path string) (segment.Segment, error)
 
+	OpenUsing(path string, config map[string]interface{}) (segment.Segment, error)
+
 	// Merge takes a set of Segments, and creates a new segment on disk at
 	// the specified path.
 	// Drops is a set of bitmaps (one for each segment) indicating which
@@ -67,6 +71,10 @@
 	Merge(segments []segment.Segment, drops []*roaring.Bitmap, path string,
 		closeCh chan struct{}, s segment.StatsReporter) (
 		[][]uint64, uint64, error)
+
+	MergeUsing(segments []segment.Segment, drops []*roaring.Bitmap, path string,
+		closeCh chan struct{}, s segment.StatsReporter, config map[string]interface{}) (
+		[][]uint64, uint64, error)
 }
 
 var supportedSegmentPlugins map[string]map[uint32]SegmentPlugin