diff --git a/centroid_index_test.go b/centroid_index_test.go new file mode 100644 index 000000000..a7334236b --- /dev/null +++ b/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) + } +} diff --git a/go.mod b/go.mod index 4736f4f88..c2ec2c4e6 100644 --- a/go.mod +++ b/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 diff --git a/index.go b/index.go index 2f1ba5fbf..41b1559c0 100644 --- a/index.go +++ b/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 +} diff --git a/index/scorch/persister.go b/index/scorch/persister.go index 977097097..6f33ed41a 100644 --- a/index/scorch/persister.go +++ b/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 { diff --git a/index/scorch/scorch.go b/index/scorch/scorch.go index efe052935..4b328ddef 100644 --- a/index/scorch/scorch.go +++ b/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) diff --git a/index/scorch/snapshot_index.go b/index/scorch/snapshot_index.go index 3422d9a14..688f9d903 100644 --- a/index/scorch/snapshot_index.go +++ b/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 } diff --git a/index/scorch/snapshot_segment.go b/index/scorch/snapshot_segment.go index 85b69c1ef..bd09287dc 100644 --- a/index/scorch/snapshot_segment.go +++ b/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 } diff --git a/index/scorch/train_noop.go b/index/scorch/train_noop.go new file mode 100644 index 000000000..3ae5ad155 --- /dev/null +++ b/index/scorch/train_noop.go @@ -0,0 +1,55 @@ +// Copyright (c) 2026 Couchbase, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build !vectors +// +build !vectors + +package scorch + +import ( + "fmt" + + index "github.com/blevesearch/bleve_index_api" + bolt "go.etcd.io/bbolt" +) + +func initTrainer(s *Scorch, config map[string]interface{}) *noopTrainer { + return &noopTrainer{} +} + +type noopTrainer struct { +} + +func (t *noopTrainer) trainLoop() {} + +func (t *noopTrainer) train(batch *index.Batch) error { + return fmt.Errorf("training is not supported with this build") +} + +func (t *noopTrainer) loadTrainedData(bucket *bolt.Bucket) error { + // noop + return nil +} + +func (t *noopTrainer) getInternal(key []byte) ([]byte, error) { + return nil, nil +} + +func (t *noopTrainer) copyFileLOCKED(file string, d index.IndexDirectory) error { + return nil +} + +func (t *noopTrainer) updateBolt(snapshotsBucket *bolt.Bucket, key []byte, value []byte) error { + return nil +} diff --git a/index/scorch/train_vector.go b/index/scorch/train_vector.go new file mode 100644 index 000000000..73dd1b029 --- /dev/null +++ b/index/scorch/train_vector.go @@ -0,0 +1,424 @@ +// Copyright (c) 2026 Couchbase, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build vectors +// +build vectors + +package scorch + +import ( + "bytes" + "encoding/binary" + "fmt" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + + "github.com/RoaringBitmap/roaring/v2" + "github.com/blevesearch/bleve/v2/mapping" + "github.com/blevesearch/bleve/v2/util" + index "github.com/blevesearch/bleve_index_api" + "github.com/blevesearch/go-faiss" + segment "github.com/blevesearch/scorch_segment_api/v2" + bolt "go.etcd.io/bbolt" +) + +type trainRequest struct { + ackCh chan error + sample segment.Segment + + // metadata to write out to bolt + trainComplete bool + vecCount uint64 + internalData []byte +} + +func initTrainer(s *Scorch, config map[string]interface{}) *vectorTrainer { + mappingBytes, ok := config["index_mapping"].([]byte) + if !ok { + // if the flag is set to something other than fastmerge, don't init + return nil + } + var im *mapping.IndexMappingImpl + err := util.UnmarshalJSON(mappingBytes, &im) + if err != nil { + // if we can't unmarshal the mapping, don't init the trainer + return nil + } + if im.VectorOptimization != index.IndexOptimizedFastMerge { + return nil + } + return &vectorTrainer{ + parent: s, + trainCh: make(chan *trainRequest), + } +} + +type vectorTrainer struct { + parent *Scorch + + m sync.Mutex + // not a searchable segment in the sense that it won't return + // the data vectors. can return centroid vectors + centroidIndex *SegmentSnapshot + trainCh chan *trainRequest +} + +func moveFile(sourcePath, destPath string) error { + // rename is supposed to be atomic on the same filesystem + err := os.Rename(sourcePath, destPath) + if err != nil { + return fmt.Errorf("error renaming file: %v", err) + } + return nil +} + +// this is not a routine that will be running throughout the lifetime of the index. It's purpose +// is to only train the vector index before the data ingestion starts. +func (t *vectorTrainer) trainLoop() { + defer func() { + t.parent.asyncTasks.Done() + }() + // initialize stuff + var totalSamplesProcessed uint64 + if t.centroidIndex != nil { + totalSamplesProcessed = t.centroidIndex.cachedMeta.fetchMeta("trainedSamples").(uint64) + } + buf := make([]byte, binary.MaxVarintLen64) + t.parent.segmentConfig[index.CentroidIndexCallback] = t.getCentroidIndex + path := filepath.Join(t.parent.path, index.CentroidIndexFileName) + for { + select { + case <-t.parent.closeCh: + return + case trainReq := <-t.trainCh: + sampleSeg := trainReq.sample + if t.centroidIndex == nil { + switch seg := sampleSeg.(type) { + case segment.UnpersistedSegment: + err := persistToDirectory(seg, nil, path) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error persisting segment: %v", err) + close(trainReq.ackCh) + return + } + default: + } + } else if sampleSeg != nil { + // merge the new segment with the existing one, to create a new + // .tmp centroid index file and then move it to the actual + // centroid index file path (during the merge, Os.Open(centroidIndexPath) + // won't be safe since its still being used for merge) + t.parent.segmentConfig[index.TrainingKey] = true + _, _, err := t.parent.segPlugin.MergeUsing([]segment.Segment{t.centroidIndex.segment, sampleSeg}, + []*roaring.Bitmap{nil, nil}, path+".tmp", t.parent.closeCh, nil, t.parent.segmentConfig) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error merging centroid index: %v", err) + close(trainReq.ackCh) + } + // reset the training flag once completed + t.parent.segmentConfig[index.TrainingKey] = false + + // close the existing centroid segment - it's supposed to be gc'd at this point + t.centroidIndex.segment.Close() + err = moveFile(path+".tmp", path) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error renaming centroid index: %v", err) + close(trainReq.ackCh) + } + } + // a bolt transaction is necessary for failover-recovery scenario and also serves as a checkpoint + // where we can be sure that the centroid index is available for the indexing operations downstream + // + // note: when the scale increases massively especially with real world dimensions of 1536+, this API + // will have to be refactored to persist in a more resource efficient way. so having this bolt related + // code will help in tracking the progress a lot better and avoid any redudant data streaming operations. + // + // todo: rethink the frequency of bolt writes + tx, err := t.parent.rootBolt.Begin(true) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error starting bolt transaction: %v", err) + close(trainReq.ackCh) + return + } + defer func() { + if err != nil { + _ = tx.Rollback() + } + }() + + snapshotsBucket, err := tx.CreateBucketIfNotExists(util.BoltSnapshotsBucket) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error creating snapshots bucket: %v", err) + close(trainReq.ackCh) + return + } + + trainerBucket, err := snapshotsBucket.CreateBucketIfNotExists(util.BoltTrainerKey) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error creating centroid bucket: %v", err) + close(trainReq.ackCh) + return + } + + // update the path of the centroid index segmentSnaphshot in bolt + err = trainerBucket.Put(util.BoltPathKey, []byte(index.CentroidIndexFileName)) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error updating centroid bucket: %v", err) + close(trainReq.ackCh) + return + } + + // train status - value is controlled by the application layer + var comp byte + if trainReq.trainComplete { + comp = 1 + } + err = trainerBucket.Put(util.BoltTrainCompleteKey, []byte{comp}) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error updating train complete bucket: %v", err) + close(trainReq.ackCh) + return + } + + totalSamplesProcessed += trainReq.vecCount + // track progress of training in terms of samples processed + binary.LittleEndian.PutUint64(buf, totalSamplesProcessed) + err = trainerBucket.Put(util.BoltTrainedSamplesKey, buf) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error updating trained samples bucket: %v", err) + close(trainReq.ackCh) + return + } + + // training related internal data that needs to be stored as per + // application layer via the SetInternal and GetInternal APIs + err = trainerBucket.Put(util.BoltInternalKey, trainReq.internalData) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error updating train internal bucket: %v", err) + close(trainReq.ackCh) + return + } + err = tx.Commit() + if err != nil { + trainReq.ackCh <- fmt.Errorf("error committing bolt transaction: %v", err) + close(trainReq.ackCh) + return + } + + err = t.parent.rootBolt.Sync() + if err != nil { + trainReq.ackCh <- fmt.Errorf("error on bolt sync: %v", err) + close(trainReq.ackCh) + return + } + + // update the centroid index pointer + centroidIndex, err := t.parent.segPlugin.OpenUsing(path, t.parent.segmentConfig) + if err != nil { + trainReq.ackCh <- fmt.Errorf("error opening centroid index: %v", err) + close(trainReq.ackCh) + return + } + t.m.Lock() + t.centroidIndex = &SegmentSnapshot{ + segment: centroidIndex, + cachedMeta: &cachedMeta{meta: nil}, + } + t.m.Unlock() + + // if the train complete flag has been set, exit the routine and cleanup + if trainReq.trainComplete { + // cleanup .tmp file if it exists + if _, err := os.Stat(path + ".tmp"); err == nil { + err = os.Remove(path + ".tmp") + if err != nil { + trainReq.ackCh <- fmt.Errorf("error removing .tmp file: %v", err) + } + } + close(trainReq.ackCh) + return + } + close(trainReq.ackCh) + } + } +} + +// loads the metadata specific to the centroid index from boltdb +func (t *vectorTrainer) loadTrainedData(bucket *bolt.Bucket) error { + if bucket == nil { + return nil + } + segmentSnapshot, err := t.parent.loadSegment(bucket) + if err != nil { + return err + } + + internalData := bucket.Get(util.BoltInternalKey) + trainedSamples := bucket.Get(util.BoltTrainedSamplesKey) + trainComplete := bucket.Get(util.BoltTrainCompleteKey) + + segmentSnapshot.cachedMeta.updateMeta("internalData", internalData) + segmentSnapshot.cachedMeta.updateMeta("trainComplete", trainComplete) + segmentSnapshot.cachedMeta.updateMeta("trainedSamples", binary.LittleEndian.Uint64(trainedSamples)) + + t.m.Lock() + defer t.m.Unlock() + t.centroidIndex = segmentSnapshot + + return nil +} + +func (t *vectorTrainer) train(batch *index.Batch) error { + // regulate the Train function + t.parent.FireIndexEvent() + + var trainData []index.Document + for key, doc := range batch.IndexOps { + if doc != nil { + // insert _id field + // no need to track updates/deletes over here since + // the API is singleton + doc.AddIDField() + } + if strings.HasPrefix(key, index.TrainDataPrefix) { + trainData = append(trainData, doc) + } + } + + var seg segment.Segment + var fin bool + var err error + trainComplete := batch.InternalOps[string(util.BoltTrainCompleteKey)] + if trainComplete == nil { + trainComplete = []byte("false") + } + fin, err = strconv.ParseBool(string(trainComplete)) + if err != nil { + return fmt.Errorf("error parsing train complete: %v", err) + } + + if !fin { + // just builds a new vector index out of the train data provided + // this is not necessarily the final train data since this is submitted + // as a request to the trainer component to be merged. once the training + // is complete, the template will be used for other operations down the line + // like merge and search. + // + // note: this might index text data too, how to handle this? s.segmentConfig? + // todo: updates/deletes -> data drift detection + seg, _, err = t.parent.segPlugin.NewUsing(trainData, t.parent.segmentConfig) + if err != nil { + return err + } + } + + trainReq := &trainRequest{ + vecCount: uint64(len(trainData)), + sample: seg, + trainComplete: fin, + ackCh: make(chan error), + } + + t.trainCh <- trainReq + err = <-trainReq.ackCh + if err != nil { + return fmt.Errorf("train_vector: train() err'd out with: %w", err) + } + + return err +} + +func (t *vectorTrainer) getInternal(key []byte) ([]byte, error) { + // todo: return the total number of vectors that have been processed so far in training + // in cbft use that as a checkpoint to resume training for n-x samples. + if t.centroidIndex != nil { + switch string(key) { + case string(util.BoltTrainCompleteKey): + trainComplete := t.centroidIndex.cachedMeta.fetchMeta("trainComplete").([]byte) + if trainComplete[0] == 1 { + return []byte("true"), nil + } + return []byte("false"), nil + case string(util.BoltTrainedSamplesKey): + // keep rv in a human readable format + trainedSamples := t.centroidIndex.cachedMeta.fetchMeta("trainedSamples").(uint64) + return []byte(fmt.Sprintf("%d", trainedSamples)), nil + // keep the default fetch to be from the internal data + default: + internalData := t.centroidIndex.cachedMeta.fetchMeta("internalData").([]byte) + return internalData, nil + } + } + return nil, nil +} + +func (t *vectorTrainer) getCentroidIndex(field string) (*faiss.IndexImpl, error) { + // return the coarse quantizer of the centroid index belonging to the field + centroidIndexSegment, ok := t.centroidIndex.segment.(segment.CentroidIndexSegment) + if !ok { + return nil, fmt.Errorf("segment is not a centroid index segment") + } + + coarseQuantizer, err := centroidIndexSegment.GetCoarseQuantizer(field) + if err != nil { + return nil, err + } + return coarseQuantizer, nil +} + +func (t *vectorTrainer) copyFileLOCKED(file string, d index.IndexDirectory) error { + if strings.HasSuffix(file, index.CentroidIndexFileName) { + // centroid index file - this is outside the snapshots domain so the bolt update is different + err := d.UpdateFileInBolt(util.BoltTrainerKey, []byte(file)) + if err != nil { + return fmt.Errorf("error updating dest index bolt: %w", err) + } + } + + return nil +} + +func (t *vectorTrainer) updateBolt(snapshotsBucket *bolt.Bucket, key []byte, value []byte) error { + if bytes.Equal(key, util.BoltTrainerKey) { + trainerBucket, err := snapshotsBucket.CreateBucketIfNotExists(util.BoltTrainerKey) + if err != nil { + return err + } + if trainerBucket == nil { + return fmt.Errorf("trainer bucket not found") + } + + // guard against duplicate updates + existingValue := trainerBucket.Get(util.BoltPathKey) + if existingValue != nil { + return fmt.Errorf("key already exists %v %v", t.parent.path, string(existingValue)) + } + + err = trainerBucket.Put(util.BoltPathKey, value) + if err != nil { + return err + } + + // update the centroid index pointer + t.centroidIndex, err = t.parent.loadSegment(trainerBucket) + if err != nil { + return err + } + } + + return nil +} diff --git a/index_alias_impl.go b/index_alias_impl.go index 451a8a910..d0895838b 100644 --- a/index_alias_impl.go +++ b/index_alias_impl.go @@ -103,6 +103,25 @@ func (i *indexAliasImpl) IndexSynonym(id string, collection string, definition * return ErrorSynonymSearchNotSupported } +func (i *indexAliasImpl) Train(batch *Batch) error { + i.mutex.RLock() + defer i.mutex.RUnlock() + + if !i.open { + return ErrorIndexClosed + } + + err := i.isAliasToSingleIndex() + if err != nil { + return err + } + + if vi, ok := i.indexes[0].(VectorIndex); ok { + return vi.Train(batch) + } + return fmt.Errorf("not a vector index") +} + func (i *indexAliasImpl) Delete(id string) error { i.mutex.RLock() defer i.mutex.RUnlock() diff --git a/index_impl.go b/index_impl.go index 586dacb3b..47026ffa6 100644 --- a/index_impl.go +++ b/index_impl.go @@ -72,9 +72,9 @@ func indexStorePath(path string) string { return path + string(os.PathSeparator) + storePath } -func newIndexUsing(path string, mapping mapping.IndexMapping, indexType string, kvstore string, kvconfig map[string]interface{}) (*indexImpl, error) { +func newIndexUsing(path string, im mapping.IndexMapping, indexType string, kvstore string, kvconfig map[string]interface{}) (*indexImpl, error) { // first validate the mapping - err := mapping.Validate() + err := im.Validate() if err != nil { return nil, err } @@ -90,7 +90,7 @@ func newIndexUsing(path string, mapping mapping.IndexMapping, indexType string, rv := indexImpl{ path: path, name: path, - m: mapping, + m: im, meta: newIndexMeta(indexType, kvstore, kvconfig), } rv.stats = &IndexStat{i: &rv} @@ -107,6 +107,12 @@ func newIndexUsing(path string, mapping mapping.IndexMapping, indexType string, kvconfig["path"] = "" } + mappingBytes, err := util.MarshalJSON(im) + if err != nil { + return nil, err + } + kvconfig["index_mapping"] = mappingBytes + // open the index indexTypeConstructor := registry.IndexTypeConstructorByName(rv.meta.IndexType) if indexTypeConstructor == nil { @@ -128,10 +134,6 @@ func newIndexUsing(path string, mapping mapping.IndexMapping, indexType string, }(&rv) // now persist the mapping - mappingBytes, err := util.MarshalJSON(mapping) - if err != nil { - return nil, err - } err = rv.i.SetInternal(util.MappingInternalKey, mappingBytes) if err != nil { return nil, err @@ -369,6 +371,20 @@ func (i *indexImpl) IndexSynonym(id string, collection string, definition *Synon return err } +func (i *indexImpl) Train(batch *Batch) error { + i.mutex.RLock() + defer i.mutex.RUnlock() + + if !i.open { + return ErrorIndexClosed + } + + if vi, ok := i.i.(index.VectorIndex); ok { + return vi.Train(batch.internal) + } + return fmt.Errorf("not a vector index") +} + // IndexAdvanced takes a document.Document object // skips the mapping and indexes it. func (i *indexImpl) IndexAdvanced(doc *document.Document) (err error) { @@ -1442,7 +1458,7 @@ func (i *indexImpl) CopyTo(d index.Directory) (err error) { err = copyReader.CopyTo(d) if err != nil { - return fmt.Errorf("error copying index metadata: %v", err) + return fmt.Errorf("error copying index data: %v", err) } // copy the metadata @@ -1577,3 +1593,35 @@ func (i *indexImpl) buildTopNCollector(ctx context.Context, req *SearchRequest, } return newCollector(), nil } + +func (i *indexImpl) CopyFile(file string, d index.IndexDirectory) (err error) { + i.mutex.RLock() + defer i.mutex.RUnlock() + + if !i.open { + return ErrorIndexClosed + } + + copyIndex, ok := i.i.(index.IndexFileCopyable) + if !ok { + return fmt.Errorf("index implementation does not support copy reader") + } + + return copyIndex.CopyFile(file, d) +} + +func (i *indexImpl) UpdateFileInBolt(key []byte, value []byte) error { + i.mutex.RLock() + defer i.mutex.RUnlock() + + if !i.open { + return ErrorIndexClosed + } + + copyIndex, ok := i.i.(index.IndexFileCopyable) + if !ok { + return fmt.Errorf("index implementation does not support file copy") + } + + return copyIndex.UpdateFileInBolt(key, value) +} diff --git a/mapping/index.go b/mapping/index.go index 143ff5a31..26a109605 100644 --- a/mapping/index.go +++ b/mapping/index.go @@ -58,6 +58,7 @@ type IndexMappingImpl struct { IndexDynamic bool `json:"index_dynamic"` DocValuesDynamic bool `json:"docvalues_dynamic"` CustomAnalysis *customAnalysis `json:"analysis,omitempty"` + VectorOptimization string `json:"vector_optimization,omitempty"` cache *registry.Cache } @@ -328,6 +329,12 @@ func (im *IndexMappingImpl) UnmarshalJSON(data []byte) error { return err } + case "vector_optimization": + err := util.UnmarshalJSON(v, &im.VectorOptimization) + if err != nil { + return err + } + default: invalidKeys = append(invalidKeys, k) } diff --git a/mapping/mapping_vectors.go b/mapping/mapping_vectors.go index 7c7ff1b98..260362f88 100644 --- a/mapping/mapping_vectors.go +++ b/mapping/mapping_vectors.go @@ -135,6 +135,20 @@ func processVector(vecI interface{}, dims int) ([]float32, bool) { return rv, true } +func (fm *FieldMapping) vectorOptimizationForPath(path []string, context *walkContext) string { + optimizationType := fm.VectorIndexOptimizedFor + if optimizationType == "" || (optimizationType == index.DefaultIndexOptimization && + context.im.VectorOptimization == index.IndexOptimizedFastMerge) { + // todo: need to support document mapping default setting as well + // if optimization type is not set at field level, or is set to default, + // and index mapping optimization type is fast merge, + // then apply fast merge optimization since it includes the default recall + // optimization as part of it. + optimizationType = context.im.VectorOptimization + } + return optimizationType +} + func (fm *FieldMapping) processVector(propertyMightBeVector interface{}, pathString string, path []string, indexes []uint64, context *walkContext) bool { vector, ok := processVector(propertyMightBeVector, fm.Dims) @@ -147,10 +161,8 @@ func (fm *FieldMapping) processVector(propertyMightBeVector interface{}, if similarity == "" { similarity = index.DefaultVectorSimilarityMetric } - vectorIndexOptimizedFor := fm.VectorIndexOptimizedFor - if vectorIndexOptimizedFor == "" { - vectorIndexOptimizedFor = index.DefaultIndexOptimization - } + vectorIndexOptimizedFor := fm.vectorOptimizationForPath(path, context) + // normalize raw vector if similarity is cosine // Since the vector can be multi-vector (flattened array of multiple vectors), // we use NormalizeMultiVector to normalize each sub-vector independently. @@ -181,10 +193,7 @@ func (fm *FieldMapping) processVectorBase64(propertyMightBeVectorBase64 interfac if similarity == "" { similarity = index.DefaultVectorSimilarityMetric } - vectorIndexOptimizedFor := fm.VectorIndexOptimizedFor - if vectorIndexOptimizedFor == "" { - vectorIndexOptimizedFor = index.DefaultIndexOptimization - } + vectorIndexOptimizedFor := fm.vectorOptimizationForPath(path, context) decodedVector, err := document.DecodeVector(encodedString) if err != nil || len(decodedVector) != fm.Dims { return diff --git a/util/keys.go b/util/keys.go index b71a7f48b..98c03f374 100644 --- a/util/keys.go +++ b/util/keys.go @@ -17,6 +17,9 @@ package util var ( // Bolt keys BoltSnapshotsBucket = []byte{'s'} + BoltTrainerKey = []byte{'t'} + BoltTrainCompleteKey = []byte{'c'} + BoltTrainedSamplesKey = []byte{'v'} BoltPathKey = []byte{'p'} BoltDeletedKey = []byte{'d'} BoltInternalKey = []byte{'i'}