add some stats and cleanup

Author: Thejas-bhat

index/scorch/stats.go

@@ -136,6 +136,9 @@ type Stats struct {
 	MaxMemMergeZapTime      uint64
 	TotMemMergeSegments     uint64
 	TotMemorySegmentsAtRoot uint64
+
+	TotTrainedSamples uint64
+	TotTrainTime      uint64
 }
 
 // atomically populates the returned map

index/scorch/train_vector.go

@@ -28,6 +28,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/blevesearch/bleve/v2/util"
@@ -132,6 +133,7 @@ func (t *vectorTrainer) trainLoop() {
 		t.parent.asyncTasks.Done()
 	}()
 
+	timeNow := time.Now()
 	path := filepath.Join(t.parent.path, index.TrainedIndexFileName)
 	for {
 		select {
@@ -215,6 +217,8 @@ func (t *vectorTrainer) trainLoop() {
 			// exit the trainer loop we've ingested the final sample set and training
 			// is assumed to be complete.
 			if t.trainingComplete.Load() {
+				atomic.StoreUint64(&t.parent.stats.TotTrainedSamples, t.trainedSamples)
+				atomic.StoreUint64(&t.parent.stats.TotTrainTime, uint64(time.Since(timeNow).Milliseconds()))
 				return
 			}
 		}
@@ -302,19 +306,16 @@ func (t *vectorTrainer) train(batch *index.Batch) error {
 }
 
 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.
 	switch string(key) {
 	case string(util.BoltTrainCompleteKey):
-		t.m.RLock()
-		defer t.m.RUnlock()
-		return []byte(fmt.Sprintf("%t", t.centroidIndex != nil)), nil
+		return []byte(strconv.FormatBool(t.trainingComplete.Load())), nil
 	}
 	return nil, nil
 }
 
 func (t *vectorTrainer) getCentroidIndex(field string) (interface{}, error) {
-	// return the coarse quantizer of the centroid index belonging to the field
+	// return the coarse quantizer of the trained faiss index belonging to the field
+	// if its not available then zap performs naive merge
 	t.m.RLock()
 	defer t.m.RUnlock()
 	if t.centroidIndex != nil {