Repository: blevesearch/bleve
Pull request #1994 (2024-03-11 07:04:19)

MB-60914: Supporting field specific metrics
 - These metrics are calculated based on analysis results initially and directly from the SegmentBase after a merge
 - This map will be stored in segment Introduction and segment snapshot
 - After merge, field level stats are collected from segment base after checking if it implements the FieldStatsReporter interface defined in scorch_segment_api
 - This map is also persisted in bolt to avoid recomputation under the key "stats" as part of the persisted part of the segment snapshot
 - FieldStats have been added to the StatsMap function of scorch. This will give an aggregated value of the stats with respect to the segments but specific to the field name and the stat using the stat name field:<field name>:<stat name>. For example, field:vector:num_vectors


Patch:
diff --git a/go.mod b/go.mod
index 490bbc24..078ac98e 100644
--- a/go.mod
+++ b/go.mod
@@ -11,7 +11,7 @@ require (
 	github.com/blevesearch/go-porterstemmer v1.0.3
 	github.com/blevesearch/goleveldb v1.0.1
 	github.com/blevesearch/gtreap v0.1.1
-	github.com/blevesearch/scorch_segment_api/v2 v2.2.8
+	github.com/blevesearch/scorch_segment_api/v2 v2.2.9
 	github.com/blevesearch/segment v0.9.1
 	github.com/blevesearch/snowball v0.6.1
 	github.com/blevesearch/snowballstem v0.9.0
@@ -23,7 +23,7 @@ require (
 	github.com/blevesearch/zapx/v13 v13.3.10
 	github.com/blevesearch/zapx/v14 v14.3.10
 	github.com/blevesearch/zapx/v15 v15.3.13
-	github.com/blevesearch/zapx/v16 v16.0.11
+	github.com/blevesearch/zapx/v16 v16.0.12
 	github.com/couchbase/moss v0.2.0
 	github.com/golang/protobuf v1.3.2
 	github.com/spf13/cobra v1.7.0
diff --git a/go.sum b/go.sum
index f612d386..2cba39cd 100644
--- a/go.sum
+++ b/go.sum
@@ -19,8 +19,8 @@ github.com/blevesearch/gtreap v0.1.1/go.mod h1:QaQyDRAT51sotthUWAH4Sj08awFSSWzgY
 github.com/blevesearch/mmap-go v1.0.2/go.mod h1:ol2qBqYaOUsGdm7aRMRrYGgPvnwLe6Y+7LMvAB5IbSA=
 github.com/blevesearch/mmap-go v1.0.4 h1:OVhDhT5B/M1HNPpYPBKIEJaD0F3Si+CrEKULGCDPWmc=
 github.com/blevesearch/mmap-go v1.0.4/go.mod h1:EWmEAOmdAS9z/pi/+Toxu99DnsbhG1TIxUoRmJw/pSs=
-github.com/blevesearch/scorch_segment_api/v2 v2.2.8 h1:+OLW38LuRKio6N6V0gIk1srwFz79FJ5v2sNqHz2HVAA=
-github.com/blevesearch/scorch_segment_api/v2 v2.2.8/go.mod h1:ckbeb7knyOOvAdZinn/ASbB7EA3HoagnJkmEV3J7+sg=
+github.com/blevesearch/scorch_segment_api/v2 v2.2.9 h1:3nBaSBRFokjE4FtPW3eUDgcAu3KphBg1GP07zy/6Uyk=
+github.com/blevesearch/scorch_segment_api/v2 v2.2.9/go.mod h1:ckbeb7knyOOvAdZinn/ASbB7EA3HoagnJkmEV3J7+sg=
 github.com/blevesearch/segment v0.9.1 h1:+dThDy+Lvgj5JMxhmOVlgFfkUtZV2kw49xax4+jTfSU=
 github.com/blevesearch/segment v0.9.1/go.mod h1:zN21iLm7+GnBHWTao9I+Au/7MBiL8pPFtJBJTsk6kQw=
 github.com/blevesearch/snowball v0.6.1 h1:cDYjn/NCH+wwt2UdehaLpr2e4BwLIjN4V/TdLsL+B5A=
@@ -43,8 +43,8 @@ github.com/blevesearch/zapx/v14 v14.3.10 h1:SG6xlsL+W6YjhX5N3aEiL/2tcWh3DO75Bnz7
 github.com/blevesearch/zapx/v14 v14.3.10/go.mod h1:qqyuR0u230jN1yMmE4FIAuCxmahRQEOehF78m6oTgns=
 github.com/blevesearch/zapx/v15 v15.3.13 h1:6EkfaZiPlAxqXz0neniq35my6S48QI94W/wyhnpDHHQ=
 github.com/blevesearch/zapx/v15 v15.3.13/go.mod h1:Turk/TNRKj9es7ZpKK95PS7f6D44Y7fAFy8F4LXQtGg=
-github.com/blevesearch/zapx/v16 v16.0.11 h1:YK7fShQ5NTPAbOnrVdxA9rvc8B7S+vrdq2Y0BDkXOco=
-github.com/blevesearch/zapx/v16 v16.0.11/go.mod h1:fN9n4RlFI8kST69yUeSSKJf/zZzZfP5bpcjmoRinKrk=
+github.com/blevesearch/zapx/v16 v16.0.12 h1:Uccxvjmn+hQ6ywQP+wIiTpdq9LnAviGoryJOmGwAo/I=
+github.com/blevesearch/zapx/v16 v16.0.12/go.mod h1:MYnOshRfSm4C4drxx1LGRI+MVFByykJ2anDY1fxdk9Q=
 github.com/couchbase/ghistogram v0.1.0 h1:b95QcQTCzjTUocDXp/uMgSNQi8oj1tGwnJ4bODWZnps=
 github.com/couchbase/ghistogram v0.1.0/go.mod h1:s1Jhy76zqfEecpNWJfWUiKZookAFaiGOEoyzgHt9i7k=
 github.com/couchbase/moss v0.2.0 h1:VCYrMzFwEryyhRSeI+/b3tRBSeTpi/8gn5Kf6dxqn+o=
diff --git a/index/scorch/introducer.go b/index/scorch/introducer.go
index cf2da1a9..2cb1398e 100644
--- a/index/scorch/introducer.go
+++ b/index/scorch/introducer.go
@@ -30,6 +30,7 @@ type segmentIntroduction struct {
 	obsoletes map[uint64]*roaring.Bitmap
 	ids       []string
 	internal  map[string][]byte
+	stats     *fieldStats
 
 	applied           chan error
 	persisted         chan error
@@ -146,6 +147,7 @@ func (s *Scorch) introduceSegment(next *segmentIntroduction) error {
 		newss := &SegmentSnapshot{
 			id:         root.segment[i].id,
 			segment:    root.segment[i].segment,
+			stats:      root.segment[i].stats,
 			cachedDocs: root.segment[i].cachedDocs,
 			cachedMeta: root.segment[i].cachedMeta,
 			creator:    root.segment[i].creator,
@@ -193,6 +195,7 @@ func (s *Scorch) introduceSegment(next *segmentIntroduction) error {
 		newSegmentSnapshot := &SegmentSnapshot{
 			id:         next.id,
 			segment:    next.data, // take ownership of next.data's ref-count
+			stats:      next.stats,
 			cachedDocs: &cachedDocs{cache: nil},
 			cachedMeta: &cachedMeta{meta: nil},
 			creator:    "introduceSegment",
@@ -281,6 +284,7 @@ func (s *Scorch) introducePersist(persist *persistIntroduction) {
 				id:         segmentSnapshot.id,
 				segment:    replacement,
 				deleted:    segmentSnapshot.deleted,
+				stats:      segmentSnapshot.stats,
 				cachedDocs: segmentSnapshot.cachedDocs,
 				cachedMeta: segmentSnapshot.cachedMeta,
 				creator:    "introducePersist",
@@ -381,6 +385,7 @@ func (s *Scorch) introduceMerge(nextMerge *segmentMerge) {
 				id:         root.segment[i].id,
 				segment:    root.segment[i].segment,
 				deleted:    root.segment[i].deleted,
+				stats:      root.segment[i].stats,
 				cachedDocs: root.segment[i].cachedDocs,
 				cachedMeta: root.segment[i].cachedMeta,
 				creator:    root.segment[i].creator,
@@ -402,7 +407,6 @@ func (s *Scorch) introduceMerge(nextMerge *segmentMerge) {
 			}
 		}
 	}
-
 	// before the newMerge introduction, need to clean the newly
 	// merged segment wrt the current root segments, hence
 	// applying the obsolete segment contents to newly merged segment
@@ -423,11 +427,17 @@ func (s *Scorch) introduceMerge(nextMerge *segmentMerge) {
 	if nextMerge.new != nil &&
 		nextMerge.new.Count() > newSegmentDeleted.GetCardinality() {
 
+		stats := newFieldStats()
+		if fsr, ok := nextMerge.new.(segment.FieldStatsReporter); ok {
+			fsr.UpdateFieldStats(stats)
+		}
+
 		// put new segment at end
 		newSnapshot.segment = append(newSnapshot.segment, &SegmentSnapshot{
 			id:         nextMerge.id,
 			segment:    nextMerge.new, // take ownership for nextMerge.new's ref-count
 			deleted:    newSegmentDeleted,
+			stats:      stats,
 			cachedDocs: &cachedDocs{cache: nil},
 			cachedMeta: &cachedMeta{meta: nil},
 			creator:    "introduceMerge",
diff --git a/index/scorch/persister.go b/index/scorch/persister.go
index 6d954e79..afd518dd 100644
--- a/index/scorch/persister.go
+++ b/index/scorch/persister.go
@@ -17,6 +17,7 @@ package scorch
 import (
 	"bytes"
 	"encoding/binary"
+	"encoding/json"
 	"fmt"
 	"io"
 	"log"
@@ -424,6 +425,7 @@ func (s *Scorch) persistSnapshotMaybeMerge(snapshot *IndexSnapshot) (
 				id:      newSegmentID,
 				segment: segment.segment,
 				deleted: nil, // nil since merging handled deletions
+				stats:   nil,
 			})
 			break
 		}
@@ -602,6 +604,18 @@ func prepareBoltSnapshot(snapshot *IndexSnapshot, tx *bolt.Tx, path string,
 				return nil, nil, err
 			}
 		}
+
+		// store segment stats
+		if segmentSnapshot.stats != nil {
+			b, err := json.Marshal(segmentSnapshot.stats.Fetch())
+			if err != nil {
+				return nil, nil, err
+			}
+			err = snapshotSegmentBucket.Put(boltStatsKey, b)
+			if err != nil {
+				return nil, nil, err
+			}
+		}
 	}
 
 	return filenames, newSegmentPaths, nil
@@ -704,6 +718,7 @@ var boltMetaDataKey = []byte{'m'}
 var boltMetaDataSegmentTypeKey = []byte("type")
 var boltMetaDataSegmentVersionKey = []byte("version")
 var boltMetaDataTimeStamp = []byte("timeStamp")
+var boltStatsKey = []byte("stats")
 var TotBytesWrittenKey = []byte("TotBytesWritten")
 
 func (s *Scorch) loadFromBolt() error {
@@ -873,6 +888,18 @@ func (s *Scorch) loadSegment(segmentBucket *bolt.Bucket) (*SegmentSnapshot, erro
 			rv.deleted = deletedBitmap
 		}
 	}
+	statBytes := segmentBucket.Get(boltStatsKey)
+	if statBytes != nil {
+		var statsMap map[string]map[string]uint64
+
+		err := json.Unmarshal(statBytes, &statsMap)
+		stats := &fieldStats{statMap: statsMap}
+		if err != nil {
+			_ = segment.Close()
+			return nil, fmt.Errorf("error reading stat bytes: %v", err)
+		}
+		rv.stats = stats
+	}
 
 	return rv, nil
 }
diff --git a/index/scorch/scorch.go b/index/scorch/scorch.go
index a9e324a3..2e6435ee 100644
--- a/index/scorch/scorch.go
+++ b/index/scorch/scorch.go
@@ -428,6 +428,8 @@ func (s *Scorch) Batch(batch *index.Batch) (err error) {
 
 	var newSegment segment.Segment
 	var bufBytes uint64
+	stats := newFieldStats()
+
 	if len(analysisResults) > 0 {
 		newSegment, bufBytes, err = s.segPlugin.New(analysisResults)
 		if err != nil {
@@ -438,11 +440,14 @@ func (s *Scorch) Batch(batch *index.Batch) (err error) {
 				segB.BytesWritten())
 		}
 		atomic.AddUint64(&s.iStats.newSegBufBytesAdded, bufBytes)
+		if fsr, ok := newSegment.(segment.FieldStatsReporter); ok {
+			fsr.UpdateFieldStats(stats)
+		}
 	} else {
 		atomic.AddUint64(&s.stats.TotBatchesEmpty, 1)
 	}
 
-	err = s.prepareSegment(newSegment, ids, batch.InternalOps, batch.PersistedCallback())
+	err = s.prepareSegment(newSegment, ids, batch.InternalOps, batch.PersistedCallback(), stats)
 	if err != nil {
 		if newSegment != nil {
 			_ = newSegment.Close()
@@ -462,7 +467,7 @@ func (s *Scorch) Batch(batch *index.Batch) (err error) {
 }
 
 func (s *Scorch) prepareSegment(newSegment segment.Segment, ids []string,
-	internalOps map[string][]byte, persistedCallback index.BatchCallback) error {
+	internalOps map[string][]byte, persistedCallback index.BatchCallback, stats *fieldStats) error {
 
 	// new introduction
 	introduction := &segmentIntroduction{
@@ -470,6 +475,7 @@ func (s *Scorch) prepareSegment(newSegment segment.Segment, ids []string,
 		data:              newSegment,
 		ids:               ids,
 		internal:          internalOps,
+		stats:             stats,
 		applied:           make(chan error),
 		persistedCallback: persistedCallback,
 	}
@@ -641,6 +647,20 @@ func (s *Scorch) StatsMap() map[string]interface{} {
 	m["num_persister_nap_merger_break"] = m["TotPersisterMergerNapBreak"]
 	m["total_compaction_written_bytes"] = m["TotFileMergeWrittenBytes"]
 
+	// calculate the aggregate of all the segment's field stats
+	aggFieldStats := newFieldStats()
+	for _, segmentSnapshot := range indexSnapshot.Segments() {
+		if segmentSnapshot.stats != nil {
+			aggFieldStats.Aggregate(segmentSnapshot.stats)
+		}
+	}
+
+	aggFieldStatsMap := aggFieldStats.Fetch()
+	for statName, stats := range aggFieldStatsMap {
+		for fieldName, val := range stats {
+			m["field:"+fieldName+":"+statName] = val
+		}
+	}
 	return m
 }
 
@@ -765,3 +785,50 @@ func parseToInteger(i interface{}) (int, error) {
 		return 0, fmt.Errorf("expects int or float64 value")
 	}
 }
+
+// Holds Zap's field level stats at a segment level
+type fieldStats struct {
+	// StatName -> FieldName -> value
+	statMap map[string]map[string]uint64
+}
+
+// Add the data into the map after checking if the statname is valid
+func (fs *fieldStats) Store(statName, fieldName string, value uint64) {
+	if _, exists := fs.statMap[statName]; !exists {
+		fs.statMap[statName] = make(map[string]uint64)
+	}
+	fs.statMap[statName][fieldName] = value
+}
+
+// Combine the given stats map with the existing map
+func (fs *fieldStats) Aggregate(stats segment.FieldStats) {
+
+	statMap := stats.Fetch()
+	if statMap == nil {
+		return
+	}
+	for statName, statMap := range statMap {
+		if _, exists := fs.statMap[statName]; !exists {
+			fs.statMap[statName] = make(map[string]uint64)
+		}
+		for fieldName, val := range statMap {
+			if _, exists := fs.statMap[statName][fieldName]; !exists {
+				fs.statMap[statName][fieldName] = 0
+			}
+			fs.statMap[statName][fieldName] += val
+		}
+	}
+}
+
+// Returns the stats map
+func (fs *fieldStats) Fetch() map[string]map[string]uint64 {
+	return fs.statMap
+}
+
+// Initializes an empty stats map
+func newFieldStats() *fieldStats {
+	rv := &fieldStats{
+		statMap: map[string]map[string]uint64{},
+	}
+	return rv
+}
diff --git a/index/scorch/snapshot_segment.go b/index/scorch/snapshot_segment.go
index f45473e0..1c14af72 100644
--- a/index/scorch/snapshot_segment.go
+++ b/index/scorch/snapshot_segment.go
@@ -39,6 +39,7 @@ type SegmentSnapshot struct {
 	segment segment.Segment
 	deleted *roaring.Bitmap
 	creator string
+	stats   *fieldStats
 
 	cachedMeta *cachedMeta
 


Test patch:
diff --git a/index/scorch/scorch_test.go b/index/scorch/scorch_test.go
index f9064b67..b8f493fe 100644
--- a/index/scorch/scorch_test.go
+++ b/index/scorch/scorch_test.go
@@ -2663,3 +2663,24 @@ func TestReadOnlyIndex(t *testing.T) {
 		t.Errorf("Expected document count to be %d got %d", 1, docCount)
 	}
 }
+
+func BenchmarkAggregateFieldStats(b *testing.B) {
+
+	fieldStatsArray := make([]*fieldStats, 1000)
+
+	for i := range fieldStatsArray {
+		fieldStatsArray[i] = newFieldStats()
+
+		fieldStatsArray[i].Store("num_vectors", "vector", uint64(rand.Intn(1000)))
+	}
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		aggFieldStats := newFieldStats()
+
+		for _, fs := range fieldStatsArray {
+			aggFieldStats.Aggregate(fs)
+		}
+	}
+}
diff --git a/search_knn_test.go b/search_knn_test.go
index 4748f843..b54ce5a9 100644
--- a/search_knn_test.go
+++ b/search_knn_test.go
@@ -1123,3 +1123,62 @@ func TestNestedVectors(t *testing.T) {
 		}
 	}
 }
+
+func TestNumVecsStat(t *testing.T) {
+
+	dataset, _, err := readDatasetAndQueries(testInputCompressedFile)
+	if err != nil {
+		t.Fatal(err)
+	}
+	documents := makeDatasetIntoDocuments(dataset)
+
+	indexMapping := NewIndexMapping()
+
+	contentFieldMapping := NewTextFieldMapping()
+	contentFieldMapping.Analyzer = en.AnalyzerName
+	indexMapping.DefaultMapping.AddFieldMappingsAt("content", contentFieldMapping)
+
+	vecFieldMapping1 := mapping.NewVectorFieldMapping()
+	vecFieldMapping1.Dims = testDatasetDims
+	vecFieldMapping1.Similarity = index.EuclideanDistance
+	indexMapping.DefaultMapping.AddFieldMappingsAt("vector", vecFieldMapping1)
+
+	tmpIndexPath := createTmpIndexPath(t)
+	index, err := New(tmpIndexPath, indexMapping)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := index.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	for i := 0; i < 10; i++ {
+		batch := index.NewBatch()
+		for j := 0; j < 3; j++ {
+			for k := 0; k < 10; k++ {
+				err := batch.Index(fmt.Sprintf("%d", i*30+j*10+k), documents[j*10+k])
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+		}
+		err = index.Batch(batch)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	statsMap := index.StatsMap()
+
+	if indexStats, exists := statsMap["index"]; exists {
+		if indexStatsMap, ok := indexStats.(map[string]interface{}); ok {
+			v1, ok := indexStatsMap["field:vector:num_vectors"].(uint64)
+			if !ok || v1 != uint64(300) {
+				t.Fatalf("mismatch in the number of vectors, expected 300, got %d", indexStatsMap["field:vector:num_vectors"])
+			}
+		}
+	}
+}