diff --git a/build.go b/build.go index 8bfa4921..3b0d353d 100644 --- a/build.go +++ b/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() diff --git a/faiss_vector_cache.go b/faiss_vector_cache.go index d5f2b294..59e0520a 100644 --- a/faiss_vector_cache.go +++ b/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, diff --git a/faiss_vector_index.go b/faiss_vector_index.go index 18c67435..acc79066 100644 --- a/faiss_vector_index.go +++ b/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 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/faiss_vector_test.go b/faiss_vector_test.go index afdd741d..c1c0502b 100644 --- a/faiss_vector_test.go +++ b/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, _, diff --git a/go.mod b/go.mod index b62a16a5..ce5eee88 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 39214f45..85530350 100644 --- a/go.sum +++ b/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= diff --git a/merge.go b/merge.go index 8526e9d1..6939b72b 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" @@ -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 diff --git a/new.go b/new.go index b79baff3..56f1c475 100644 --- a/new.go +++ b/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,6 +301,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 +309,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 +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 { diff --git a/section_faiss_vector_index.go b/section_faiss_vector_index.go index 4b8ffe80..3530f9f5 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,20 @@ 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() { + elapsed := uint64(time.Since(start)) + atomic.AddUint64(&vo.stats.TotVecSectionVecsProcessedTime, elapsed) + if vo.trainingPhase { + atomic.AddUint64(&vo.stats.TotVecSectionTrainingPhaseVecsProcessedTime, elapsed) + } + }() + if fieldID == math.MaxUint16 { // doc processing checkpoint - no action needed return @@ -866,6 +915,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 +933,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 +943,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 +961,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(seg.Stats) + } return rv } @@ -936,6 +990,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 *seg.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 +1000,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 +1038,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.(*seg.Stats) } } @@ -992,9 +1062,11 @@ type faissIndexConfig struct { optimizationType string nlist int useGPU bool + stats *seg.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 *seg.Stats) *faissIndexConfig { return &faissIndexConfig{ indexType: idxType, dimension: dimension, @@ -1003,12 +1075,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/segment.go b/segment.go index 157102e7..62343290 100644 --- a/segment.go +++ b/segment.go @@ -49,15 +49,23 @@ 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) { + zapStats, ok := config[segment.StatsKey].(*segment.Stats) + if !ok || zapStats == nil { + zapStats = new(segment.Stats) + } + + atomic.AddUint64(&zapStats.TotOpenBeg, 1) f, err := os.Open(path) if err != nil { + atomic.AddUint64(&zapStats.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(&zapStats.TotOpenErrors, 1) return nil, err } @@ -73,6 +81,7 @@ func (*ZapPlugin) open(path string, config map[string]interface{}) (segment.Segm trainedIndexCache: newTrainedIndexCache(), fieldDvReaders: make([][]*docValueReader, len(segmentSections)), config: config, + stats: zapStats, }, f: f, mm: mm, @@ -84,18 +93,21 @@ func (*ZapPlugin) open(path string, config map[string]interface{}) (segment.Segm err = rv.loadConfig() if err != nil { _ = rv.Close() + atomic.AddUint64(&zapStats.TotOpenErrors, 1) return nil, err } err = rv.loadFields() if err != nil { _ = rv.Close() + atomic.AddUint64(&zapStats.TotOpenErrors, 1) return nil, err } err = rv.loadDvReaders() if err != nil { _ = rv.Close() + atomic.AddUint64(&zapStats.TotOpenErrors, 1) return nil, err } @@ -103,9 +115,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(&zapStats.TotOpenErrors, 1) return nil, err } + atomic.AddUint64(&zapStats.TotOpenEnd, 1) return rv, nil } @@ -144,6 +158,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 *segment.Stats } func (sb *SegmentBase) Size() int { @@ -712,6 +729,7 @@ func (s *Segment) closeActual() (err error) { } } + atomic.AddUint64(&s.stats.TotSegmentsClosed, 1) return }