matchtree: capture Stats before pruning (#607)

keegancsmith · web-flow · commit 93f7b0c983a6 · 2023-06-29T12:17:43.000+02:00
We now call updateStats upto twice per shard search. The intention of
this is to capture statistics before pruning the matchtree. Previously
we would of done work in creating a matchtree but would then prune those
items away and would then never capture those statistics.

In practice that work was reading just one or two varint (the size of a
posting list) so likely had minimal impact on the reported statistics.
However, in the next commit we want to introduce a statistic which is
recorded even if we generate a noMatchTree.

The main technical part of this commit is ensuring all existing
updateStats functions can be called twice without overcounting.

Test Plan: go test
diff --git a/eval.go b/eval.go
@@ -181,6 +181,9 @@ func (d *indexData) Search(ctx context.Context, q query.Q, opts *SearchOptions)
 		return nil, err
 	}
 
+	// Capture the costs of construction before pruning
+	updateMatchTreeStats(mt, &res.Stats)
+
 	mt, err = pruneMatchTree(mt)
 	if err != nil {
 		return nil, err
@@ -380,11 +383,8 @@ nextFileMatch:
 		}
 	}
 
-	visitMatchTree(mt, func(mt matchTree) {
-		if atom, ok := mt.(interface{ updateStats(*Stats) }); ok {
-			atom.updateStats(&res.Stats)
-		}
-	})
+	// Update stats based on work done during document search.
+	updateMatchTreeStats(mt, &res.Stats)
 
 	// If document ranking is enabled, then we can rank and truncate the files to save memory.
 	if limit := opts.MaxDocDisplayCount; opts.UseDocumentRanks && limit > 0 && limit < len(res.Files) {
diff --git a/hititer.go b/hititer.go
@@ -183,18 +183,19 @@ func (i *inMemoryIterator) next(limit uint32) {
 // compressedPostingIterator goes over a delta varint encoded posting
 // list.
 type compressedPostingIterator struct {
-	blob, orig []byte
-	_first     uint32
-	what       ngram
+	blob             []byte
+	indexBytesLoaded int
+	_first           uint32
+	what             ngram
 }
 
 func newCompressedPostingIterator(b []byte, w ngram) *compressedPostingIterator {
 	d, sz := binary.Uvarint(b)
 	return &compressedPostingIterator{
-		_first: uint32(d),
-		blob:   b[sz:],
-		orig:   b,
-		what:   w,
+		_first:           uint32(d),
+		blob:             b[sz:],
+		indexBytesLoaded: sz,
+		what:             w,
 	}
 }
 
@@ -216,6 +217,7 @@ func (i *compressedPostingIterator) next(limit uint32) {
 	for i._first <= limit && len(i.blob) > 0 {
 		delta, sz := binary.Uvarint(i.blob)
 		i._first += uint32(delta)
+		i.indexBytesLoaded += sz
 		i.blob = i.blob[sz:]
 	}
 
@@ -225,7 +227,8 @@ func (i *compressedPostingIterator) next(limit uint32) {
 }
 
 func (i *compressedPostingIterator) updateStats(s *Stats) {
-	s.IndexBytesLoaded += int64(len(i.orig) - len(i.blob))
+	s.IndexBytesLoaded += int64(i.indexBytesLoaded)
+	i.indexBytesLoaded = 0
 }
 
 // mergingIterator forms the merge of a set of hitIterators, to
diff --git a/matchiter.go b/matchiter.go
@@ -63,6 +63,10 @@ type matchIterator interface {
 	docIterator
 
 	candidates() []*candidateMatch
+
+	// updateStats is called twice. After matchtree construction and after
+	// searching is done. Implementations must take care to not report
+	// statistics twice.
 	updateStats(*Stats)
 }
 
@@ -150,6 +154,7 @@ func (i *ngramDocIterator) prepare(nextDoc uint32) {
 func (i *ngramDocIterator) updateStats(s *Stats) {
 	i.iter.updateStats(s)
 	s.NgramMatches += i.matchCount
+	i.matchCount = 0
 }
 
 func (i *ngramDocIterator) candidates() []*candidateMatch {
diff --git a/matchtree.go b/matchtree.go
@@ -520,6 +520,16 @@ func visitMatchTree(t matchTree, f func(matchTree)) {
 	}
 }
 
+// updateMatchTreeStats calls updateStats on all atoms in mt which have that
+// function defined.
+func updateMatchTreeStats(mt matchTree, stats *Stats) {
+	visitMatchTree(mt, func(mt matchTree) {
+		if atom, ok := mt.(interface{ updateStats(*Stats) }); ok {
+			atom.updateStats(stats)
+		}
+	})
+}
+
 // visitMatches visits all atoms which can contribute matches. Note: This
 // skips noVisitMatchTree.
 func visitMatches(t matchTree, known map[matchTree]bool, f func(matchTree)) {
