MB-72082: add more stats around vector search

build.go

@@ -20,6 +20,7 @@ import (
 	"io"
 	"math"
 	"os"
+	"sync/atomic"
 
 	"github.com/RoaringBitmap/roaring/v2"
 	index "github.com/blevesearch/bleve_index_api"
@@ -32,7 +33,14 @@ const Type string = "zap"
 const fieldNotUninverted uint64 = math.MaxUint64
 
 func (sb *SegmentBase) Persist(path string) error {
-	return PersistSegmentBase(sb, path)
+	atomic.AddUint64(&sb.stats.TotPersistBeg, 1)
+	err := PersistSegmentBase(sb, path)
+	if err != nil {
+		atomic.AddUint64(&sb.stats.TotPersistErrors, 1)
+		return err
+	}
+	atomic.AddUint64(&sb.stats.TotPersistEnd, 1)
+	return nil
 }
 
 // WriteTo is an implementation of io.WriterTo interface.
@@ -98,8 +106,10 @@ func PersistSegmentBase(sb *SegmentBase, path string) error {
 func rewriteSegmentBase(sb *SegmentBase, path string) error {
 	closeCh := make(chan struct{})
 	defer close(closeCh)
+	config := map[string]interface{}{statsKey: sb.stats}
+
 	_, _, err := mergeSegmentBases([]*SegmentBase{sb}, []*roaring.Bitmap{nil},
-		path, DefaultChunkMode, closeCh, nil, nil)
+		path, DefaultChunkMode, closeCh, nil, config)
 	if err != nil {
 		return err
 	}
@@ -207,6 +217,14 @@ func InitSegmentBase(mem []byte, memCRC uint32, chunkMode uint32, numDocs uint64
 		fieldsInv:     make([]string, 0),
 		config:        config,
 	}
+	// extract stats from config if present, otherwise allocate a throwaway
+	// instance so all increment sites remain nil-check-free.
+	sb.stats = new(Stats)
+	if config != nil {
+		if s, ok := config[statsKey].(*Stats); ok && s != nil {
+			sb.stats = s
+		}
+	}
 	sb.updateSize()
 
 	// initialize the file reader with an empty callback

faiss_vector_cache.go

@@ -72,6 +72,7 @@ type vectorCacheOptions struct {
 	optStr  string
 
 	skipMapping bool // if true, skip building the idMapping
+	stats       *Stats
 }
 
 func newVectorCacheOptions(mem []byte, numDocs uint32, except *roaring.Bitmap,

faiss_vector_index.go

@@ -60,6 +60,7 @@ type faissIndexParams struct {
 	nlist int
 	// ioFlags used to read the index from bytes
 	ioFlags int
+	stats   *Stats
 }
 
 // newFaissIndexParams constructs a faissIndexParams with the given optimization

faiss_vector_index_float32.go

@@ -21,6 +21,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"reflect"
+	"sync/atomic"
 
 	index "github.com/blevesearch/bleve_index_api"
 	faiss "github.com/blevesearch/go-faiss"
@@ -112,7 +113,7 @@ func (f *faissFloat32Index) write(buf []byte, w *FileWriter) error {
 		return err
 	}
 	idxBytes = w.process(idxBytes)
-
+	atomic.AddUint64(&f.params.stats.TotVecSectionFloatIndexBytesWritten, uint64(len(idxBytes)))
 	// write the length of the serialized vector index bytes
 	n := binary.PutUvarint(buf, uint64(len(idxBytes)))
 	_, err = w.Write(buf[:n])

merge.go

@@ -22,6 +22,7 @@ import (
 	"math"
 	"os"
 	"sort"
+	"sync/atomic"
 
 	"github.com/RoaringBitmap/roaring/v2"
 	index "github.com/blevesearch/bleve_index_api"
@@ -48,7 +49,7 @@ func (z *ZapPlugin) MergeUsing(segments []seg.Segment, drops []*roaring.Bitmap,
 	return z.merge(segments, drops, path, closeCh, s, config)
 }
 
-func (*ZapPlugin) merge(segments []seg.Segment, drops []*roaring.Bitmap, path string,
+func (z *ZapPlugin) merge(segments []seg.Segment, drops []*roaring.Bitmap, path string,
 	closeCh chan struct{}, s seg.StatsReporter, config map[string]interface{}) (
 	[][]uint64, uint64, error) {
 	segmentBases := make([]*SegmentBase, len(segments))
@@ -62,7 +63,28 @@ func (*ZapPlugin) merge(segments []seg.Segment, drops []*roaring.Bitmap, path st
 			panic(fmt.Sprintf("oops, unexpected segment type: %T", segment))
 		}
 	}
-	return mergeSegmentBases(segmentBases, drops, path, DefaultChunkMode, closeCh, s, config)
+
+	config[statsKey] = &z.stats
+
+	atomic.AddUint64(&z.stats.TotMergesBeg, 1)
+	atomic.AddUint64(&z.stats.TotMergeInputSegments, uint64(len(segments)))
+	var totalInputDocs, droppedDocs uint64
+	for i, sb := range segmentBases {
+		totalInputDocs += sb.numDocs
+		if drops[i] != nil {
+			droppedDocs += drops[i].GetCardinality()
+		}
+	}
+	atomic.AddUint64(&z.stats.TotMergeDroppedDocs, droppedDocs)
+	atomic.AddUint64(&z.stats.TotMergeOutputDocs, totalInputDocs-droppedDocs)
+
+	newDocNums, size, err := mergeSegmentBases(segmentBases, drops, path, DefaultChunkMode, closeCh, s, config)
+	if err != nil {
+		atomic.AddUint64(&z.stats.TotMergesErrors, 1)
+		return nil, 0, err
+	}
+	atomic.AddUint64(&z.stats.TotMergesEnd, 1)
+	return newDocNums, size, nil
 }
 
 func mergeSegmentBases(segmentBases []*SegmentBase, drops []*roaring.Bitmap, path string,
@@ -205,6 +227,7 @@ func mergeToWriter(segments []*SegmentBase, drops []*roaring.Bitmap,
 		"fieldsMap":     fieldsMap,
 		"numDocs":       numDocs,
 		"fieldsOptions": fieldsOptions,
+		"stats":         config[statsKey].(*Stats),
 	}
 	if config != nil {
 		args["config"] = config

new.go

@@ -50,10 +50,11 @@ func (z *ZapPlugin) NewUsing(results []index.Document, config map[string]interfa
 	return z.newWithChunkMode(results, DefaultChunkMode, config)
 }
 
-func (*ZapPlugin) newWithChunkMode(results []index.Document,
+func (z *ZapPlugin) newWithChunkMode(results []index.Document,
 	chunkMode uint32, config map[string]interface{}) (segment.Segment, uint64, error) {
 	s := interimPool.Get().(*interim)
 
+	s.stats = &z.stats
 	var br bytes.Buffer
 	if s.lastNumDocs > 0 {
 		// use previous results to initialize the buf with an estimate
@@ -67,6 +68,7 @@ func (*ZapPlugin) newWithChunkMode(results []index.Document,
 		br.Grow(estimateAvgBytesPerDoc * estimateNumResults)
 	}
 
+	atomic.AddUint64(&s.stats.TotNewRootDocsProcessed, uint64(len(results)))
 	var err error
 	s.results, s.edgeList = flattenNestedDocuments(results, s.edgeList)
 	s.config = config
@@ -81,6 +83,10 @@ func (*ZapPlugin) newWithChunkMode(results []index.Document,
 
 	sb, err := InitSegmentBase(br.Bytes(), s.w.Sum32(), chunkMode,
 		uint64(len(s.results)), storedIndexOffset, sectionsIndexOffset, config)
+	if err == nil {
+		// propagate stats to the SegmentBase so that Persist() can track flush stats
+		sb.stats = s.stats
+	}
 
 	// get the bytes written before the interim's reset() call
 	// write it to the newly formed segment base.
@@ -134,6 +140,8 @@ type interim struct {
 	lastOutSize int
 
 	opaque map[int]resetable
+
+	stats *Stats
 }
 
 func (s *interim) reset() (err error) {
@@ -225,6 +233,7 @@ func (s *interim) convert() (uint64, uint64, error) {
 		"fieldsMap":     s.FieldsMap,
 		"fieldsInv":     s.FieldsInv,
 		"fieldsOptions": s.FieldsOptions,
+		"stats":         s.stats,
 	}
 	if s.config != nil {
 		args["config"] = s.config
@@ -291,14 +300,17 @@ func (s *interim) processDocuments() {
 	for docNum, result := range s.results {
 		s.processDocument(uint32(docNum), result)
 	}
+	atomic.AddUint64(&s.stats.TotNewDocsProcessed, uint64(len(s.results)))
 }
 
 func (s *interim) processDocument(docNum uint32,
 	result index.Document) {
 	// this callback is essentially going to be invoked on each field,
 	// as part of which preprocessing, cumulation etc. of the doc's data
 	// will take place.
+	var fieldCount int
 	visitField := func(field index.Field) {
+		fieldCount++
 		fieldID := uint16(s.getOrDefineField(field.Name()))
 
 		// section specific processing of the field
@@ -324,6 +336,11 @@ func (s *interim) processDocument(docNum uint32,
 		section.Process(s.opaque, docNum, nil, math.MaxUint16)
 	}
 
+	if fieldCount > 0 {
+		atomic.AddUint64(&s.stats.TotNewDocsIndexed, 1)
+	} else {
+		atomic.AddUint64(&s.stats.TotNewDocsDropped, 1)
+	}
 }
 
 func (s *interim) getBytesWritten() uint64 {

plugin.go

@@ -16,7 +16,14 @@ package zap
 
 // ZapPlugin implements the Plugin interface of
 // the blevesearch/scorch_segment_api pkg
-type ZapPlugin struct{}
+type ZapPlugin struct {
+	// ensures that a stats instance is always available for a plugin
+	stats Stats
+}
+
+func InitPlugin() *ZapPlugin {
+	return &ZapPlugin{}
+}
 
 func (*ZapPlugin) Type() string {
 	return Type