MB-62182: Improve recoverability of the index's trained state

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

@@ -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

index.go

@@ -353,6 +353,13 @@ type IndexCopyable interface {
 	CopyTo(d index.Directory) error
 }
 
+// IndexFileCopyable is an index supporting the transfer of a single file between
+// two indexes
+type IndexFileCopyable interface {
+	UpdateFileInBolt(key []byte, value []byte) error
+	CopyFile(file string, d index.IndexDirectory) error
+}
+
 // FileSystemDirectory is the default implementation for the
 // index.Directory interface.
 type FileSystemDirectory string
@@ -396,3 +403,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/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{
@@ -575,6 +574,11 @@ func copyToDirectory(srcPath string, d index.Directory) (int64, error) {
 		return 0, fmt.Errorf("GetWriter err: %v", err)
 	}
 
+	// skip
+	if dest == nil {
+		return 0, nil
+	}
+
 	sourceFileStat, err := os.Stat(srcPath)
 	if err != nil {
 		return 0, err
@@ -853,6 +857,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 {
@@ -861,6 +869,7 @@ func (s *Scorch) loadFromBolt() error {
 		if snapshots == nil {
 			return nil
 		}
+		var mappingBytes []byte
 		foundRoot := false
 		c := snapshots.Cursor()
 		for k, _ := c.Last(); k != nil; k, _ = c.Prev() {
@@ -902,8 +911,21 @@ func (s *Scorch) loadFromBolt() error {
 				_ = rootPrev.DecRef()
 			}
 
+			mappingBytes, err = indexSnapshot.GetInternal(util.MappingInternalKey)
 			foundRoot = true
 		}
+
+		// try init trainer with the mapping details
+		if s.trainer == nil {
+			s.config["index_mapping"] = mappingBytes
+			s.trainer = initTrainer(s, s.config)
+		}
+		trainerBucket := snapshots.Bucket(util.BoltTrainerKey)
+		err := s.trainer.loadTrainedData(trainerBucket)
+		if err != nil {
+			return err
+		}
+
 		return nil
 	})
 	if err != nil {

index/scorch/scorch.go

@@ -17,6 +17,7 @@ package scorch
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"sync"
@@ -78,6 +79,10 @@ type Scorch struct {
 	persisterNotifier        chan *epochWatcher
 	rootBolt                 *bolt.DB
 	asyncTasks               sync.WaitGroup
+	// not a real searchable segment
+	centroidIndex *SegmentSnapshot
+
+	trainer trainer
 
 	onEvent      func(event Event) bool
 	onAsyncError func(err error, path string)
@@ -89,6 +94,33 @@ 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)
+
+	// file transfer operations
+	copyFileLOCKED(file string, d index.IndexDirectory) error
+	updateBolt(snapshotsBucket *bolt.Bucket, key []byte, value []byte) error
+}
+
 type ScorchErrorType string
 
 func (t ScorchErrorType) Error() string {
@@ -170,6 +202,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 {
@@ -207,6 +244,10 @@ func NewScorch(storeName string,
 		return nil, err
 	}
 
+	if trainer := initTrainer(rv, config); trainer != nil {
+		rv.trainer = trainer
+	}
+
 	return rv, nil
 }
 
@@ -261,6 +302,11 @@ func (s *Scorch) Open() error {
 	s.asyncTasks.Add(1)
 	go s.introducerLoop()
 
+	if s.trainer != nil {
+		s.asyncTasks.Add(1)
+		go s.trainer.trainLoop()
+	}
+
 	if !s.readOnly && s.path != "" {
 		s.asyncTasks.Add(1)
 		go s.persisterLoop()
@@ -534,6 +580,31 @@ 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):
+		fallthrough
+	case string(util.BoltTrainedSamplesKey):
+		if s.trainer != nil {
+			return s.trainer.getInternal(key)
+		} else {
+			return nil, fmt.Errorf("getInternal on train keys is not supported" +
+				" with this build")
+		}
+	}
+	return nil, nil
+}
+
+func (s *Scorch) Train(batch *index.Batch) error {
+	if s.trainer != nil {
+		return s.trainer.train(batch)
+	}
+	return fmt.Errorf("training is not supported with this build")
+}
+
 func (s *Scorch) prepareSegment(newSegment segment.Segment, ids []string,
 	internalOps map[string][]byte, persistedCallback index.BatchCallback, stats *fieldStats,
 ) error {
@@ -979,6 +1050,65 @@ func (s *Scorch) CopyReader() index.CopyReader {
 	return rv
 }
 
+func (s *Scorch) UpdateFileInBolt(key []byte, value []byte) error {
+	tx, err := s.rootBolt.Begin(true)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err != nil {
+			_ = tx.Rollback()
+		}
+	}()
+
+	snapshotsBucket, err := tx.CreateBucketIfNotExists(util.BoltSnapshotsBucket)
+	if err != nil {
+		return err
+	}
+
+	// currently this is specific to centroid index file update
+	err = s.trainer.updateBolt(snapshotsBucket, key, value)
+	if err != nil {
+		return err
+	}
+
+	err = tx.Commit()
+	if err != nil {
+		return err
+	}
+
+	return s.rootBolt.Sync()
+}
+
+// CopyFile copies a specific file to a destination directory which has an access to a bleve index
+// doing a io.Copy() isn't enough because the file needs to be tracked in bolt file as well
+func (s *Scorch) CopyFile(file string, d index.IndexDirectory) error {
+	s.rootLock.Lock()
+	defer s.rootLock.Unlock()
+
+	dest, err := d.GetWriter(filepath.Join("store", file))
+	if err != nil {
+		return err
+	}
+
+	source, err := os.Open(filepath.Join(s.path, file))
+	if err != nil {
+		return err
+	}
+
+	defer source.Close()
+	defer dest.Close()
+	_, err = io.Copy(dest, source)
+	if err != nil {
+		return err
+	}
+
+	// this code is currently specific to copying trained data but is future proofed for other files
+	// to be updated in the dest's bolt
+	err = s.trainer.copyFileLOCKED(file, d)
+	return err
+}
+
 // external API to fire a scorch event (EventKindIndexStart) externally from bleve
 func (s *Scorch) FireIndexEvent() {
 	s.fireEvent(EventKindIndexStart, 0)

index/scorch/snapshot_index.go

@@ -466,6 +466,10 @@ func (is *IndexSnapshot) Fields() ([]string, error) {
 }
 
 func (is *IndexSnapshot) GetInternal(key []byte) ([]byte, error) {
+	_, ok := is.internal[string(key)]
+	if !ok {
+		return is.parent.getInternal(key)
+	}
 	return is.internal[string(key)], nil
 }
 

index/scorch/snapshot_segment.go

@@ -370,8 +370,11 @@ func (c *cachedMeta) updateMeta(field string, val interface{}) {
 
 func (c *cachedMeta) fetchMeta(field string) (rv interface{}) {
 	c.m.RLock()
+	defer c.m.RUnlock()
+	if c.meta == nil {
+		return nil
+	}
 	rv = c.meta[field]
-	c.m.RUnlock()
 	return rv
 }