diff --git a/build.go b/build.go index 8bfa4921..d5c5c5b8 100644 --- a/build.go +++ b/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 diff --git a/faiss_vector_cache.go b/faiss_vector_cache.go index d5f2b294..5fbf13f3 100644 --- a/faiss_vector_cache.go +++ b/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, diff --git a/faiss_vector_index.go b/faiss_vector_index.go index 18c67435..cf5a40c9 100644 --- a/faiss_vector_index.go +++ b/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 diff --git a/faiss_vector_index_float32.go b/faiss_vector_index_float32.go index 9220168e..79abefd1 100644 --- a/faiss_vector_index_float32.go +++ b/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]) diff --git a/merge.go b/merge.go index 8526e9d1..0d79c8c5 100644 --- a/merge.go +++ b/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 diff --git a/new.go b/new.go index b79baff3..2249f5fe 100644 --- a/new.go +++ b/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,6 +300,7 @@ 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, @@ -298,7 +308,9 @@ func (s *interim) processDocument(docNum uint32, // 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 { diff --git a/plugin.go b/plugin.go index f67297ec..0c1272e8 100644 --- a/plugin.go +++ b/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 diff --git a/section_faiss_vector_index.go b/section_faiss_vector_index.go index 4b8ffe80..e9638be0 100644 --- a/section_faiss_vector_index.go +++ b/section_faiss_vector_index.go @@ -23,6 +23,7 @@ import ( "fmt" "math" "sync/atomic" + "time" "github.com/RoaringBitmap/roaring/v2" index "github.com/blevesearch/bleve_index_api" @@ -117,6 +118,8 @@ func (v *faissVectorIndexSection) Merge(opaque map[int]resetable, segments []*Se drops []*roaring.Bitmap, fieldsInv []string, newDocNumsIn [][]uint64, w *FileWriter, closeCh chan struct{}) error { vo := v.getVectorIndexOpaque(opaque) + + var totalVecFields int // preallocating the space over here, if there are too many fields // in the segment this will help by avoiding multiple allocation // calls. @@ -126,6 +129,7 @@ func (v *faissVectorIndexSection) Merge(opaque map[int]resetable, segments []*Se indexes := make([]*vecIndexInfo, 0, len(segments)) // mapping from vector IDs to docIDs across segments vecToDocID := make([]uint64, 0, len(segments)) + // for every field, gather the vector indexes from the segments // that have them, merge them and write them out to the writer. for fieldID, fieldName := range fieldsInv { @@ -153,7 +157,6 @@ func (v *faissVectorIndexSection) Merge(opaque map[int]resetable, segments []*Se if pos == 0 { continue } - // loading doc values - adhering to the sections format. never // valid values for vector section _, n := binary.Uvarint(sb.mem[pos : pos+binary.MaxVarintLen64]) @@ -212,6 +215,8 @@ func (v *faissVectorIndexSection) Merge(opaque map[int]resetable, segments []*Se continue } + atomic.AddUint64(&vo.stats.TotVecSectionVecsDeleted, uint64(newIndexInfo.nvecs-len(newIndexInfo.vecIds))) + // read the type of vector index indexType, n := binary.Uvarint(sb.mem[pos : pos+binary.MaxVarintLen64]) pos += n @@ -232,10 +237,13 @@ func (v *faissVectorIndexSection) Merge(opaque map[int]resetable, segments []*Se continue } + totalVecFields++ + count := w.Count() err := vo.flushSectionMetadata(fieldID, w, vecToDocID, indexes) if err != nil { return err } + atomic.AddUint64(&vo.stats.TotVecSectionMetadataBytesWritten, uint64(w.Count()-count)) // we're going to use the trained index template regardless of whether there's // a update/delete in the segments being merged and we let the fast merge @@ -251,6 +259,10 @@ func (v *faissVectorIndexSection) Merge(opaque map[int]resetable, segments []*Se return err } } + + if totalVecFields > int(atomic.LoadUint64(&vo.stats.TotVecSectionFieldsIndexed)) { + atomic.StoreUint64(&vo.stats.TotVecSectionFieldsIndexed, uint64(totalVecFields)) + } return nil } @@ -436,6 +448,7 @@ func (v *vectorIndexOpaque) fastMergeIndexes(trainedIndex faissIndexIVF, cfg *fa if err != nil { return err } + atomic.AddUint64(&v.stats.TotVecSectionVecsReconstructed, uint64(len(vi.vecIds))) } else { if err = ivfMergedIdx.mergeFrom(childIdx, mergedIdx.ntotal()); err != nil { // either the childIdx isn't compatible for fast merge or merge_from failed @@ -445,6 +458,9 @@ func (v *vectorIndexOpaque) fastMergeIndexes(trainedIndex faissIndexIVF, cfg *fa if err != nil { return err } + atomic.AddUint64(&v.stats.TotVecSectionVecsReconstructed, uint64(len(vi.vecIds))) + } else { + atomic.AddUint64(&v.stats.TotVecSectionFastMerges, 1) } } } @@ -469,6 +485,14 @@ func (v *vectorIndexOpaque) mergeAndWriteVectorIndexes(trainedIndex faissIndexIV var indexType faissIndexType var validMerge bool + atomic.AddUint64(&v.stats.TotVecSectionMergesBegin, 1) + start := time.Now() + + defer func() { + atomic.AddUint64(&v.stats.TotVecSectionMergeTime, uint64(time.Since(start))) + atomic.AddUint64(&v.stats.TotVecSectionMergesEnd, 1) + }() + for segI, segBase := range sbs { // Considering merge operations on vector indexes are expensive, it is // worth including an early exit if the merge is aborted, saving us @@ -497,7 +521,7 @@ func (v *vectorIndexOpaque) mergeAndWriteVectorIndexes(trainedIndex faissIndexIV ioFlags = faissIOFlagsReadOnly } reconsParams := newFaissIndexParams(currVecIndex.indexOptimizedFor, currVecIndex.nvecs, 0, ioFlags) - + reconsParams.stats = v.stats // load binary index from disk if present if currVecIndex.indexType == faissBIVFIndex { // get to the bivf part of the vector index section @@ -513,7 +537,6 @@ func (v *vectorIndexOpaque) mergeAndWriteVectorIndexes(trainedIndex faissIndexIV vecIndexes[segI].index, err = newFaissBinaryIndexFromBytes(bIndexBytes, fIndexBytes, reconsParams) } else { vecIndexes[segI].index, err = newFaissFloat32IndexFromBytes(fIndexBytes, reconsParams) - } if err != nil { freeReconstructedIndexes(vecIndexes) @@ -554,9 +577,10 @@ func (v *vectorIndexOpaque) mergeAndWriteVectorIndexes(trainedIndex faissIndexIV // We perform fast merge whenever a compatible trained index is available, // regardless of whether the GPU is enabled for this field. if canFastMerge(trainedIndex, indexOptimizedFor, nvecs) { - config := newFaissIndexConfig(indexType, indexOptimizedFor, dims, metric, nvecs, nlist, false) + config := newFaissIndexConfig(indexType, indexOptimizedFor, dims, metric, nvecs, nlist, false, v.stats) err := v.fastMergeIndexes(trainedIndex, config, drops, vecIndexes, w, closeCh) if err != nil { + atomic.AddUint64(&v.stats.TotVecSectionFastMergeErrs, 1) return err } // free the indexes as we won't need them anymore after the fast merge @@ -565,8 +589,9 @@ func (v *vectorIndexOpaque) mergeAndWriteVectorIndexes(trainedIndex faissIndexIV } // Reconstruct Merge Path: - config := newFaissIndexConfig(indexType, indexOptimizedFor, dims, metric, nvecs, nlist, useGPU) + config := newFaissIndexConfig(indexType, indexOptimizedFor, dims, metric, nvecs, nlist, useGPU, v.stats) // merging of indexes with reconstruction method. + atomic.AddUint64(&v.stats.TotVecSectionNaiveMerges, 1) // the vecIds in each index contain only the valid vectors, // so we reconstruct only those. indexData := make([]float32, 0, indexDataCap) @@ -590,6 +615,7 @@ func (v *vectorIndexOpaque) mergeAndWriteVectorIndexes(trainedIndex faissIndexIV freeReconstructedIndexes(vecIndexes) return err } + atomic.AddUint64(&v.stats.TotVecSectionVecsReconstructed, uint64(currNumVecs)) indexData = append(indexData, recons...) } } @@ -626,13 +652,23 @@ func (v *vectorIndexOpaque) writeFaissIndex(vecs *vectorSet, config *faissIndexC // and nprobe. The order matters for GPU indexes: CloneToCPU (done inside // trainAndAdd) clears the direct map and nprobe, so they must be set after. if ivfIndex := index.castIVF(); ivfIndex != nil { + atomic.AddUint64(&v.stats.TotVecSectionIVFIndexesCreated, 1) // train the vector index and add the vectors to it. The training step // performs k-means clustering to partition the data space such that during // search time we probe only a subset of vectors (non-exhaustive search). + start := time.Now() err = ivfIndex.trainAndAdd(vecs, vecs) if err != nil { return err } + + if v.trainingPhase { + atomic.AddUint64(&v.stats.TotVecSectionTrainingPhaseTrainingTime, uint64(time.Since(start))) + } else { + atomic.AddUint64(&v.stats.TotVecSectionTrainingTime, uint64(time.Since(start))) + } + atomic.AddUint64(&v.stats.TotVecSectionTrainOps, 1) + // the direct map maintained in the IVF index is essential for the // reconstruction of vectors based on the sequential vector IDs in the // future merges use direct map type 1 -> array based direct map, since @@ -645,6 +681,7 @@ func (v *vectorIndexOpaque) writeFaissIndex(vecs *vectorSet, config *faissIndexC nprobe := calculateNprobe(config.nlist, config.optimizationType) ivfIndex.setNProbe(nprobe) } else { + atomic.AddUint64(&v.stats.TotVecSectionFlatIndexesCreated, 1) // add the vectors to the index using sequential vector IDs starting // from 0 to N-1 err = index.add(vecs) @@ -773,6 +810,7 @@ func (vo *vectorIndexOpaque) writeVectorIndexes(w *FileWriter) error { // d. index optimization type // e. vectorID -> docID mapping tempBuf := vo.grabBuf(binary.MaxVarintLen64) + start := time.Now() for fieldID, content := range vo.fieldVectorIndex { // number of vectors to be indexed for this field nvecs := len(content.vecDocIDs) @@ -836,13 +874,14 @@ func (vo *vectorIndexOpaque) writeVectorIndexes(w *FileWriter) error { if err != nil { return err } + atomic.AddUint64(&vo.stats.TotVecSectionMetadataBytesWritten, uint64(w.Count()-fieldStart)) nlist := vo.numCentroids(nvecs) // determine the type of vector index to be created based on the index optimization // and create the faiss index for the vectors associated with this field and // write out the index into the segment writer. indexType := determineIndexTypeFromOptimization(content.optimizedFor) - config := newFaissIndexConfig(indexType, content.optimizedFor, content.dimension, metric, nvecs, nlist, content.useGPU) + config := newFaissIndexConfig(indexType, content.optimizedFor, content.dimension, metric, nvecs, nlist, content.useGPU, vo.stats) err = vo.writeFaissIndex(vecSet, config, w) if err != nil { return err @@ -852,10 +891,19 @@ func (vo *vectorIndexOpaque) writeVectorIndexes(w *FileWriter) error { vo.incrementBytesWritten(uint64(w.Count() - fieldStart)) vo.fieldAddrs[fieldID] = fieldStart } + atomic.AddUint64(&vo.stats.TotVecSectionIndexWriteTime, uint64(time.Since(start))) return nil } func (vo *vectorIndexOpaque) process(field index.VectorField, fieldID uint16, docNum uint32) { + start := time.Now() + defer func() { + atomic.AddUint64(&vo.stats.TotVecSectionVecsProcessedTime, uint64(time.Since(start))) + if vo.trainingPhase { + atomic.AddUint64(&vo.stats.TotVecSectionTrainingPhaseVecsProcessedTime, uint64(time.Since(start))) + } + }() + if fieldID == math.MaxUint16 { // doc processing checkpoint - no action needed return @@ -866,6 +914,7 @@ func (vo *vectorIndexOpaque) process(field index.VectorField, fieldID uint16, do metric := field.Similarity() indexOptimizedFor := field.IndexOptimizedFor() useGPU := vo.fieldsOptions[name].UseGPU() + // caller is supposed to make sure len(vec) is a multiple of dim. // Not double checking it here to avoid the overhead. // This accounts for multi-vector fields, where a field can have @@ -883,8 +932,8 @@ func (vo *vectorIndexOpaque) process(field index.VectorField, fieldID uint16, do dimension: dim, metric: metric, optimizedFor: indexOptimizedFor, - vectors: make([]float32, 0, dim*numVectors), - vecDocIDs: make([]uint32, 0, numVectors), + vectors: make([]float32, 0, dim*numVectors*vo.numDocs), + vecDocIDs: make([]uint32, 0, numVectors*vo.numDocs), useGPU: useGPU, } vo.fieldVectorIndex[fieldID] = content @@ -893,6 +942,7 @@ func (vo *vectorIndexOpaque) process(field index.VectorField, fieldID uint16, do content.vectors = append(content.vectors, vector...) content.vecDocIDs = append(content.vecDocIDs, docNum) } + atomic.AddUint64(&vo.stats.TotNewVectorsProcessed, uint64(numVectors)) } func (v *faissVectorIndexSection) getVectorIndexOpaque(opaque map[int]resetable) *vectorIndexOpaque { @@ -910,6 +960,9 @@ func (v *faissVectorIndexSection) InitOpaque(args map[string]interface{}) reseta for k, v := range args { rv.Set(k, v) } + if rv.stats == nil { + rv.stats = new(Stats) + } return rv } @@ -936,6 +989,8 @@ type vectorIndexOpaque struct { config map[string]interface{} // number of bytes written out for the vector index section, used for metrics and tracking bytesWritten uint64 + // stats holds the statistics for the vector index processing + stats *Stats // fieldAddrs maps fieldID to the address of its vector section fieldAddrs map[uint16]int // fieldVectorIndex maps fieldID to its vector index content @@ -944,6 +999,10 @@ type vectorIndexOpaque struct { fieldsOptions map[string]index.FieldIndexingOptions // tmp0 is a reusable buffer tmp0 []byte + // numDocs tracks the total number of documents processed during introduction, helpful while + // preallocating buffers for faster copy operations + numDocs int + trainingPhase bool } func (vo *vectorIndexOpaque) incrementBytesWritten(val uint64) { @@ -978,6 +1037,16 @@ func (v *vectorIndexOpaque) Set(key string, val interface{}) { v.fieldsOptions = val.(map[string]index.FieldIndexingOptions) case "config": v.config = val.(map[string]interface{}) + if v.config != nil { + if tp, ok := v.config[index.TrainingKey].(*index.TrainingParams); ok && tp != nil { + v.trainingPhase = true + } + } + + case "results": + v.numDocs = len(val.([]index.Document)) + case "stats": + v.stats = val.(*Stats) } } @@ -992,9 +1061,11 @@ type faissIndexConfig struct { optimizationType string nlist int useGPU bool + stats *Stats } -func newFaissIndexConfig(idxType faissIndexType, optimizationType string, dimension, metricType, numVecs, nlist int, useGPU bool) *faissIndexConfig { +func newFaissIndexConfig(idxType faissIndexType, optimizationType string, dimension, + metricType, numVecs, nlist int, useGPU bool, stats *Stats) *faissIndexConfig { return &faissIndexConfig{ indexType: idxType, dimension: dimension, @@ -1003,12 +1074,14 @@ func newFaissIndexConfig(idxType faissIndexType, optimizationType string, dimens nlist: nlist, optimizationType: optimizationType, useGPU: useGPU, + stats: stats, } } // Factory function to create a faissIndex for the given index config. func faissIndexFactory(cfg *faissIndexConfig) (faissIndex, error) { params := newFaissIndexParams(cfg.optimizationType, cfg.numVecs, cfg.nlist, faissIOFlags) + params.stats = cfg.stats switch cfg.indexType { case faissFP32Index: description := determineFloat32IndexToUse(cfg.numVecs, cfg.nlist, cfg.optimizationType) diff --git a/section_inverted_text_index.go b/section_inverted_text_index.go index a2a82de8..c4be173f 100644 --- a/section_inverted_text_index.go +++ b/section_inverted_text_index.go @@ -1045,6 +1045,8 @@ type invertedIndexOpaque struct { fieldsSame bool numDocs uint64 + + stats *Stats } func (io *invertedIndexOpaque) Reset() (err error) { @@ -1097,6 +1099,7 @@ func (io *invertedIndexOpaque) Reset() (err error) { io.fieldsSame = false io.numDocs = 0 + io.stats = nil clear(io.fieldAddrs) return err diff --git a/segment.go b/segment.go index 157102e7..d022ff03 100644 --- a/segment.go +++ b/segment.go @@ -49,15 +49,18 @@ func (z *ZapPlugin) Open(path string) (segment.Segment, error) { return z.open(path, nil) } -func (*ZapPlugin) open(path string, config map[string]interface{}) (segment.Segment, error) { +func (z *ZapPlugin) open(path string, config map[string]interface{}) (segment.Segment, error) { + atomic.AddUint64(&z.stats.TotOpenBeg, 1) f, err := os.Open(path) if err != nil { + atomic.AddUint64(&z.stats.TotOpenErrors, 1) return nil, err } mm, err := mmap.Map(f, mmap.RDONLY, 0) if err != nil { // mmap failed, try to close the file _ = f.Close() + atomic.AddUint64(&z.stats.TotOpenErrors, 1) return nil, err } @@ -73,6 +76,7 @@ func (*ZapPlugin) open(path string, config map[string]interface{}) (segment.Segm trainedIndexCache: newTrainedIndexCache(), fieldDvReaders: make([][]*docValueReader, len(segmentSections)), config: config, + stats: &z.stats, }, f: f, mm: mm, @@ -84,18 +88,21 @@ func (*ZapPlugin) open(path string, config map[string]interface{}) (segment.Segm err = rv.loadConfig() if err != nil { _ = rv.Close() + atomic.AddUint64(&z.stats.TotOpenErrors, 1) return nil, err } err = rv.loadFields() if err != nil { _ = rv.Close() + atomic.AddUint64(&z.stats.TotOpenErrors, 1) return nil, err } err = rv.loadDvReaders() if err != nil { _ = rv.Close() + atomic.AddUint64(&z.stats.TotOpenErrors, 1) return nil, err } @@ -103,9 +110,11 @@ func (*ZapPlugin) open(path string, config map[string]interface{}) (segment.Segm err = rv.nstIndexCache.initialize(rv.numDocs, rv.getEdgeListOffset(), rv.mem) if err != nil { _ = rv.Close() + atomic.AddUint64(&z.stats.TotOpenErrors, 1) return nil, err } + atomic.AddUint64(&z.stats.TotOpenEnd, 1) return rv, nil } @@ -144,6 +153,9 @@ type SegmentBase struct { synIndexCache *synonymIndexCache geoIndexCache *geoIndexCache nstIndexCache *nestedIndexCache + + // segment level stats that are tracked and reported as part of the segment's lifecycle + stats *Stats } func (sb *SegmentBase) Size() int { @@ -712,6 +724,7 @@ func (s *Segment) closeActual() (err error) { } } + atomic.AddUint64(&s.stats.TotSegmentsClosed, 1) return } diff --git a/stats.go b/stats.go new file mode 100644 index 00000000..05bfae0b --- /dev/null +++ b/stats.go @@ -0,0 +1,87 @@ +// 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. + +package zap + +import ( + "reflect" + "sync/atomic" + "unsafe" +) + +const statsKey = "_zap_stats" + +// all the stats we want to track for a segment during its lifetime, these are +// all uint64 to allow us to use light atomic operations +type Stats struct { + TotNewRootDocsProcessed uint64 + TotNewDocsProcessed uint64 + TotNewDocsIndexed uint64 + TotNewDocsDropped uint64 + TotNewVectorsProcessed uint64 + + TotPersistBeg uint64 + TotPersistEnd uint64 + TotPersistErrors uint64 + + TotMergesBeg uint64 + TotMergesEnd uint64 + TotMergesErrors uint64 + TotMergeInputSegments uint64 + TotMergeOutputDocs uint64 + TotMergeDroppedDocs uint64 + + TotVecSectionMergesBegin uint64 + TotVecSectionMergesEnd uint64 + TotVecSectionMergeTime uint64 + TotVecSectionVecsReconstructed uint64 + TotVecSectionIVFIndexesCreated uint64 + TotVecSectionFlatIndexesCreated uint64 + TotVecSectionTrainingTime uint64 + TotVecSectionFastMerges uint64 + TotVecSectionFastMergeErrs uint64 + TotVecSectionNaiveMerges uint64 + TotVecSectionMetadataBytesWritten uint64 + TotVecSectionFloatIndexBytesWritten uint64 + TotVecSectionVecsDeleted uint64 + TotVecSectionFieldsIndexed uint64 + TotVecSectionTrainOps uint64 + TotVecSectionIndexWriteTime uint64 + TotVecSectionVecsProcessedTime uint64 + + TotVecSectionTrainingPhaseVecsProcessedTime uint64 + TotVecSectionTrainingPhaseTrainingTime uint64 + + TotOpenBeg uint64 + TotOpenEnd uint64 + TotOpenErrors uint64 + + TotSegmentsClosed uint64 +} + +func (z *ZapPlugin) StatsMap() map[string]interface{} { + svet := reflect.TypeOf(z.stats) + n := svet.NumField() + m := make(map[string]interface{}, n) + base := unsafe.Pointer(&z.stats) + for i := 0; i < n; i++ { + field := svet.Field(i) + + // use unsafe.Pointer to avoid heap allocs, safe to do this here since all the stats + // enforced to be uint64 + p := (*uint64)(unsafe.Pointer(uintptr(base) + field.Offset)) + m[field.Name] = atomic.LoadUint64(p) + } + return m +}