MB-72082: add more stats around vector search

build.go

@@ -20,9 +20,11 @@ import (
 	"io"
 	"math"
 	"os"
+	"sync/atomic"
 
 	"github.com/RoaringBitmap/roaring/v2"
 	index "github.com/blevesearch/bleve_index_api"
+	segment "github.com/blevesearch/scorch_segment_api/v2"
 )
 
 const Version uint32 = 17
@@ -32,7 +34,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.
@@ -99,7 +108,7 @@ func rewriteSegmentBase(sb *SegmentBase, path string) error {
 	closeCh := make(chan struct{})
 	defer close(closeCh)
 	_, _, err := mergeSegmentBases([]*SegmentBase{sb}, []*roaring.Bitmap{nil},
-		path, DefaultChunkMode, closeCh, nil, nil)
+		path, DefaultChunkMode, closeCh, nil, nil, sb.stats)
 	if err != nil {
 		return err
 	}
@@ -186,7 +195,7 @@ func persistStoredFieldValues(fieldID int,
 
 func InitSegmentBase(mem []byte, memCRC uint32, chunkMode uint32, numDocs uint64,
 	storedIndexOffset uint64, sectionsIndexOffset uint64,
-	config map[string]interface{}) (*SegmentBase, error) {
+	config map[string]interface{}, stats *segment.Stats) (*SegmentBase, error) {
 	sb := &SegmentBase{
 		mem:                 mem,
 		memCRC:              memCRC,
@@ -201,11 +210,13 @@ func InitSegmentBase(mem []byte, memCRC uint32, chunkMode uint32, numDocs uint64
 		synIndexCache:       newSynonymIndexCache(),
 		geoIndexCache:       newGeoIndexCache(),
 		nstIndexCache:       newNestedIndexCache(),
+		trainedIndexCache:   newTrainedIndexCache(),
 		// following fields gets populated by loadFields
 		fieldsMap:     make(map[string]uint16),
 		fieldsOptions: make(map[string]index.FieldIndexingOptions),
 		fieldsInv:     make([]string, 0),
 		config:        config,
+		stats:         stats,
 	}
 	sb.updateSize()
 

faiss_vector_cache.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/RoaringBitmap/roaring/v2"
+	seg "github.com/blevesearch/scorch_segment_api/v2"
 )
 
 // -----------------------------------------------------------------------------
@@ -72,6 +73,7 @@ type vectorCacheOptions struct {
 	optStr  string
 
 	skipMapping bool // if true, skip building the idMapping
+	stats       *seg.Stats
 }
 
 func newVectorCacheOptions(mem []byte, numDocs uint32, except *roaring.Bitmap,

faiss_vector_index.go

@@ -23,6 +23,7 @@ import (
 	"reflect"
 
 	"github.com/blevesearch/go-faiss"
+	seg "github.com/blevesearch/scorch_segment_api/v2"
 )
 
 var (
@@ -60,6 +61,7 @@ type faissIndexParams struct {
 	nlist int
 	// ioFlags used to read the index from bytes
 	ioFlags int
+	stats   *seg.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])

faiss_vector_test.go

@@ -481,6 +481,13 @@ func TestVectorSegment(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	defer func() {
+		cerr := segOnDisk.Close()
+		if cerr != nil {
+			t.Fatalf("error closing segment on disk: %v", cerr)
+		}
+	}()
+
 	fieldsSectionsMap := vecSegBase.fieldsSectionsMap
 	stubVecFieldStartAddr := fieldsSectionsMap[vecSegBase.fieldsMap["stubVec"]-1][SectionFaissVectorIndex]
 	docValueStart, docValueEnd, indexBytesLen, _,

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/blevesearch/bleve_index_api v1.4.1-0.20260729060817-8e56340f2a7e
 	github.com/blevesearch/go-faiss v1.1.5
 	github.com/blevesearch/mmap-go v1.2.0
-	github.com/blevesearch/scorch_segment_api/v2 v2.4.9-0.20260729090843-4313bda09bee
+	github.com/blevesearch/scorch_segment_api/v2 v2.4.9
 	github.com/blevesearch/vellum v1.2.0
 	github.com/golang/snappy v1.0.0
 	github.com/spf13/cobra v1.10.2

go.sum

@@ -8,8 +8,8 @@ github.com/blevesearch/go-faiss v1.1.5 h1:/IU5lkOahH9Ghfk9n3F6N0XD7PYVXZJWmNDc9T
 github.com/blevesearch/go-faiss v1.1.5/go.mod h1:w3W9AiWsFRGVaMG+/cmJi7iHEAuGyC6blsgO1EzCK/M=
 github.com/blevesearch/mmap-go v1.2.0 h1:l33nNKPFcBjJUMwem6sAYJPUzhUCABoK9FxZDGiFNBI=
 github.com/blevesearch/mmap-go v1.2.0/go.mod h1:Vd6+20GBhEdwJnU1Xohgt88XCD/CTWcqbCNxkZpyBo0=
-github.com/blevesearch/scorch_segment_api/v2 v2.4.9-0.20260729090843-4313bda09bee h1:RmX8uCxp4RDbsW0PRFsFED5f7X5xUvz6/9cYb9HrNTU=
-github.com/blevesearch/scorch_segment_api/v2 v2.4.9-0.20260729090843-4313bda09bee/go.mod h1:WUUkAocbkDlNK/kgAE13NvS9oxe+u618mYZ8sOvcCc4=
+github.com/blevesearch/scorch_segment_api/v2 v2.4.9 h1:q1shQvpwVGp0IISmX08Ead5XFOPNmZrB761Snpru3vk=
+github.com/blevesearch/scorch_segment_api/v2 v2.4.9/go.mod h1:WUUkAocbkDlNK/kgAE13NvS9oxe+u618mYZ8sOvcCc4=
 github.com/blevesearch/vellum v1.2.0 h1:xkDiOEsHc2t3Cp0NsNZZ36pvc130sCzcGKOPMzXe+e0=
 github.com/blevesearch/vellum v1.2.0/go.mod h1:uEcfBJz7mAOf0Kvq6qoEKQQkLODBF46SINYNkZNae4k=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=

merge.go

@@ -22,6 +22,7 @@ import (
 	"math"
 	"os"
 	"sort"
+	"sync/atomic"
 
 	"github.com/RoaringBitmap/roaring/v2"
 	index "github.com/blevesearch/bleve_index_api"
@@ -62,11 +63,36 @@ 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)
+
+	zapStats, ok := config[seg.StatsKey].(*seg.Stats)
+	if !ok || zapStats == nil {
+		zapStats = new(seg.Stats)
+	}
+
+	atomic.AddUint64(&zapStats.TotMergesBeg, 1)
+	atomic.AddUint64(&zapStats.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(&zapStats.TotMergeDroppedDocs, droppedDocs)
+	atomic.AddUint64(&zapStats.TotMergeOutputDocs, totalInputDocs-droppedDocs)
+
+	newDocNums, size, err := mergeSegmentBases(segmentBases, drops, path, DefaultChunkMode, closeCh, s, config, zapStats)
+	if err != nil {
+		atomic.AddUint64(&zapStats.TotMergesErrors, 1)
+		return nil, 0, err
+	}
+	atomic.AddUint64(&zapStats.TotMergesEnd, 1)
+	return newDocNums, size, nil
 }
 
 func mergeSegmentBases(segmentBases []*SegmentBase, drops []*roaring.Bitmap, path string,
-	chunkMode uint32, closeCh chan struct{}, s seg.StatsReporter, config map[string]interface{}) (
+	chunkMode uint32, closeCh chan struct{}, s seg.StatsReporter, config map[string]interface{},
+	stats *seg.Stats) (
 	[][]uint64, uint64, error) {
 	flag := os.O_RDWR | os.O_CREATE
 
@@ -92,7 +118,7 @@ func mergeSegmentBases(segmentBases []*SegmentBase, drops []*roaring.Bitmap, pat
 	}
 
 	newDocNums, numDocs, storedIndexOffset, _, _, sectionsIndexOffset, err :=
-		mergeToWriter(segmentBases, drops, chunkMode, w, closeCh, config)
+		mergeToWriter(segmentBases, drops, chunkMode, w, closeCh, config, stats)
 	if err != nil {
 		cleanup()
 		return nil, 0, err
@@ -171,7 +197,8 @@ func finalizeFieldOptions(fieldOptions map[string]index.FieldIndexingOptions,
 }
 
 func mergeToWriter(segments []*SegmentBase, drops []*roaring.Bitmap,
-	chunkMode uint32, w *FileWriter, closeCh chan struct{}, config map[string]interface{}) (
+	chunkMode uint32, w *FileWriter, closeCh chan struct{}, config map[string]interface{},
+	stats *seg.Stats) (
 	newDocNums [][]uint64, numDocs, storedIndexOffset uint64,
 	fieldsInv []string, fieldsMap map[string]uint16, sectionsIndexOffset uint64,
 	err error) {
@@ -205,6 +232,7 @@ func mergeToWriter(segments []*SegmentBase, drops []*roaring.Bitmap,
 		"fieldsMap":     fieldsMap,
 		"numDocs":       numDocs,
 		"fieldsOptions": fieldsOptions,
+		"stats":         stats,
 	}
 	if config != nil {
 		args["config"] = config

new.go

@@ -50,10 +50,16 @@ 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)
 
+	zapStats, ok := config[segment.StatsKey].(*segment.Stats)
+	if !ok || zapStats == nil {
+		zapStats = new(segment.Stats)
+	}
+
+	s.stats = zapStats
 	var br bytes.Buffer
 	if s.lastNumDocs > 0 {
 		// use previous results to initialize the buf with an estimate
@@ -67,6 +73,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
@@ -80,7 +87,7 @@ func (*ZapPlugin) newWithChunkMode(results []index.Document,
 	}
 
 	sb, err := InitSegmentBase(br.Bytes(), s.w.Sum32(), chunkMode,
-		uint64(len(s.results)), storedIndexOffset, sectionsIndexOffset, config)
+		uint64(len(s.results)), storedIndexOffset, sectionsIndexOffset, config, s.stats)
 
 	// get the bytes written before the interim's reset() call
 	// write it to the newly formed segment base.
@@ -134,6 +141,8 @@ type interim struct {
 	lastOutSize int
 
 	opaque map[int]resetable
+
+	stats *segment.Stats
 }
 
 func (s *interim) reset() (err error) {
@@ -225,6 +234,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 +301,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 +337,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 {