wip: batch training + interfaces to reuse pre-trained file

Author: Thejas-bhat

index.go

@@ -353,6 +353,11 @@ type IndexCopyable interface {
 	CopyTo(d index.Directory) error
 }
 
+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

index/scorch/persister.go

@@ -564,11 +564,6 @@ 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
@@ -847,10 +842,34 @@ func zapFileName(epoch uint64) string {
 	return fmt.Sprintf("%012x.zap", epoch)
 }
 
+func (s *Scorch) updateCentroidIndex(bucket *bolt.Bucket) error {
+	if bucket == nil {
+		return nil
+	}
+	segmentSnapshot, err := s.loadSegment(bucket)
+	if err != nil {
+		return err
+	}
+	s.rootLock.Lock()
+	defer s.rootLock.Unlock()
+
+	s.centroidIndex = segmentSnapshot
+	return nil
+}
+
 // bolt snapshot code
 
 func (s *Scorch) loadFromBolt() error {
 	err := s.rootBolt.View(func(tx *bolt.Tx) error {
+		centroidIndexBucket := tx.Bucket(util.BoltCentroidIndexKey)
+		if centroidIndexBucket == nil {
+			return nil
+		}
+		err := s.updateCentroidIndex(centroidIndexBucket)
+		if err != nil {
+			return err
+		}
+
 		snapshots := tx.Bucket(util.BoltSnapshotsBucket)
 		if snapshots == nil {
 			return nil

index/scorch/scorch.go

@@ -15,8 +15,10 @@
 package scorch
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
@@ -513,7 +515,19 @@ 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()
+	if string(key) == "_centroid_index_complete" {
+		return []byte(fmt.Sprintf("%t", s.centroidIndex != nil)), nil
+	}
+	return nil, nil
+}
+
+// min 39 per centroid, recommeded 50
+// max 256
 func (s *Scorch) Train(batch *index.Batch) error {
+	// is the lock really needed?
 	s.rootLock.Lock()
 	defer s.rootLock.Unlock()
 	if s.centroidIndex != nil {
@@ -523,6 +537,12 @@ func (s *Scorch) Train(batch *index.Batch) error {
 	var trainData []index.Document
 	if s.centroidIndex == nil {
 		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)
 			}
@@ -537,11 +557,16 @@ func (s *Scorch) Train(batch *index.Batch) error {
 	//
 	// note: this might index text data too, how to handle this? s.segmentConfig?
 	// todo: updates/deletes -> data drift detection
-	seg, _, err := s.segPlugin.NewEx(trainData, s.segmentConfig)
+	s.segmentConfig["training"] = true
+	seg, n, err := s.segPlugin.NewEx(trainData, s.segmentConfig)
 	if err != nil {
 		return err
 	}
-	filename := "centroid_index.zap"
+	// reset the training flag once completed
+	s.segmentConfig["training"] = false
+	// not suffixing with .zap since the current garbage collection is tailored to .zap ext files
+	// we don't want to gc this file ever.
+	filename := "centroid_index"
 	path := filepath.Join(s.path, filename)
 
 	switch seg := seg.(type) {
@@ -562,7 +587,56 @@ func (s *Scorch) Train(batch *index.Batch) error {
 	s.centroidIndex = &SegmentSnapshot{
 		segment: persistedSegment,
 	}
-	s.segmentConfig["getCentroidIndexCallback"] = s.getCentroidIndex
+
+	fmt.Println("number of bytes written to centroid index", n)
+	// s.segmentConfig["getCentroidIndexCallback"] = s.getCentroidIndex
+	// updateBolt(tx, cetntroid)
+	// filename := "centroid_index"
+	// path := filepath.Join(s.path, filename)
+	// f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0600)
+	// if err != nil {
+	// 	return err
+	// }
+
+	// bufw := bufio.NewWriter(f)
+	// _, err = bufw.Write([]byte(strings.Join([]string{"centroid_index1", path}, " ")))
+	// if err != nil {
+	// 	return err
+	// }
+	// err = bufw.Flush()
+	// if err != nil {
+	// 	return err
+	// }
+	// err = f.Sync()
+	// if err != nil {
+	// 	return err
+	// }
+	// err = f.Close()
+	// if err != nil {
+	// 	return err
+	// }
+
+	tx, err := s.rootBolt.Begin(true)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	snapshotsBucket, err := tx.CreateBucketIfNotExists(util.BoltSnapshotsBucket)
+	if err != nil {
+		return err
+	}
+
+	err = snapshotsBucket.Put(util.BoltCentroidIndexKey, []byte(path))
+	if err != nil {
+		return err
+	}
+
+	err = tx.Commit()
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -1018,6 +1092,91 @@ func (s *Scorch) CopyReader() index.CopyReader {
 	return rv
 }
 
+func (s *Scorch) updateCentroidIndexInBolt(tx *bolt.Tx) error {
+	centroidIndexBucket, err := tx.CreateBucketIfNotExists(util.BoltCentroidIndexKey)
+	if err != nil {
+		return err
+	}
+
+	err = centroidIndexBucket.Put(util.BoltPathKey, []byte("centroid_index.zap"))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Scorch) UpdateFileInBolt(key []byte, value []byte) error {
+	tx, err := s.rootBolt.Begin(true)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	snapshotsBucket, err := tx.CreateBucketIfNotExists(util.BoltSnapshotsBucket)
+	if err != nil {
+		return err
+	}
+
+	// currently this is specific to centroid index file update
+	if bytes.Equal(key, util.BoltCentroidIndexKey) {
+		// guard against duplicate updates
+		existingValue := snapshotsBucket.Get(key)
+		if existingValue != nil {
+			return fmt.Errorf("key already exists")
+		}
+
+		err = snapshotsBucket.Put(key, value)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = tx.Commit()
+	if err != nil {
+		return err
+	}
+
+	err = s.rootBolt.Sync()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// 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()
+
+	// this code is currently specific to centroid index file but is future proofed for other files
+	// to be updated in the dest's bolt
+	if strings.HasSuffix(file, "centroid_index") {
+		// centroid index file - this is outside the snapshots domain so the bolt update is different
+		err := d.UpdateFileInBolt(util.BoltCentroidIndexKey, []byte(file))
+		if err != nil {
+			return err
+		}
+	}
+
+	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)
+	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

@@ -469,6 +469,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_alias_impl.go

@@ -110,14 +110,13 @@ func (i *indexAliasImpl) Train(batch *Batch) error {
 	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)
+	if vi, ok := i.indexes[0].(index.VectorIndex); ok {
+		return vi.Train(batch.internal)
 	}
 	return fmt.Errorf("not a vector index")
 }

index_impl.go

@@ -377,8 +377,8 @@ func (i *indexImpl) Train(batch *Batch) error {
 		return ErrorIndexClosed
 	}
 
-	if vi, ok := i.i.(VectorIndex); ok {
-		return vi.Train(batch)
+	if vi, ok := i.i.(index.VectorIndex); ok {
+		return vi.Train(batch.internal)
 	}
 	return fmt.Errorf("not a vector index")
 }
@@ -1376,6 +1376,38 @@ func (m *searchHitSorter) Less(i, j int) bool {
 	return c < 0
 }
 
+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)
+}
+
 // CopyTo (index.Directory, filter)
 func (i *indexImpl) CopyTo(d index.Directory) (err error) {
 	i.mutex.RLock()
@@ -1405,7 +1437,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

util/keys.go

@@ -17,6 +17,7 @@ package util
 var (
 	// Bolt keys
 	BoltSnapshotsBucket           = []byte{'s'}
+	BoltCentroidIndexKey          = []byte{'c'}
 	BoltPathKey                   = []byte{'p'}
 	BoltDeletedKey                = []byte{'d'}
 	BoltInternalKey               = []byte{'i'}