Returning bytesRead per query as part of searchResult (#142)

* refactoring for bytes read to be returned per query

* - fixing the doc values' bytesRead tracking
  by avoid the double counting of stuff
- fixing the double counting of a segment's
  loading portion of bytes read

* scraped the atomics operations for bytesRead tracking.

* code cleanup

* updating go.mod and go.sum

* bug fix: implementing BytesRead for stubDocument

Author: Thejas-bhat

dict.go

@@ -30,6 +30,8 @@ type Dictionary struct {
 	fieldID   uint16
 	fst       *vellum.FST
 	fstReader *vellum.Reader
+
+	bytesRead uint64
 }
 
 // represents an immutable, empty dictionary
@@ -125,6 +127,22 @@ func (d *Dictionary) AutomatonIterator(a segment.Automaton,
 	return emptyDictionaryIterator
 }
 
+func (d *Dictionary) incrementBytesRead(val uint64) {
+	d.bytesRead += val
+}
+
+func (d *Dictionary) BytesRead() uint64 {
+	return d.bytesRead
+}
+
+func (d *Dictionary) ResetBytesRead(val uint64) {
+	d.bytesRead = val
+}
+
+func (d *Dictionary) BytesWritten() uint64 {
+	return 0
+}
+
 // DictionaryIterator is an iterator for term dictionary
 type DictionaryIterator struct {
 	d         *Dictionary

doc_test.go

@@ -12,6 +12,10 @@ type stubDocument struct {
 	composite []*stubField
 }
 
+func (s *stubDocument) StoredFieldsBytes() uint64 {
+	return 0
+}
+
 func (s *stubDocument) ID() string {
 	return s.id
 }

docvalues.go

@@ -21,7 +21,6 @@ import (
 	"math"
 	"reflect"
 	"sort"
-	"sync/atomic"
 
 	index "github.com/blevesearch/bleve_index_api"
 	segment "github.com/blevesearch/scorch_segment_api/v2"
@@ -40,6 +39,31 @@ type docNumTermsVisitor func(docNum uint64, terms []byte) error
 type docVisitState struct {
 	dvrs    map[uint16]*docValueReader
 	segment *SegmentBase
+
+	bytesRead uint64
+}
+
+// Implements the segment.DiskStatsReporter interface
+// The purpose of this implementation is to get
+// the bytes read from the disk (pertaining to the
+// docvalues) while querying.
+// the loadDvChunk retrieves the next chunk of docvalues
+// and the bytes retrieved off the disk pertaining to that
+// is accounted as well.
+func (d *docVisitState) incrementBytesRead(val uint64) {
+	d.bytesRead += val
+}
+
+func (d *docVisitState) BytesRead() uint64 {
+	return d.bytesRead
+}
+
+func (d *docVisitState) BytesWritten() uint64 {
+	return 0
+}
+
+func (d *docVisitState) ResetBytesRead(val uint64) {
+	d.bytesRead = val
 }
 
 type docValueReader struct {
@@ -131,27 +155,12 @@ func (s *SegmentBase) loadFieldDocValueReader(field string,
 	return fdvIter, nil
 }
 
-// Implements the segment.DiskStatsReporter interface
-// The purpose of this implementation is to get
-// the bytes read from the disk (pertaining to the
-// docvalues) while querying.
-// the loadDvChunk retrieves the next chunk of docvalues
-// and the bytes retrieved off the disk pertaining to that
-// is accounted as well.
-func (di *docValueReader) BytesRead() uint64 {
-	return atomic.LoadUint64(&di.bytesRead)
-}
-
-func (di *docValueReader) ResetBytesRead(val uint64) {
-	atomic.StoreUint64(&di.bytesRead, val)
+func (d *docValueReader) getBytesRead() uint64 {
+	return d.bytesRead
 }
 
-func (di *docValueReader) incrementBytesRead(val uint64) {
-	atomic.AddUint64(&di.bytesRead, val)
-}
-
-func (di *docValueReader) BytesWritten() uint64 {
-	return 0
+func (d *docValueReader) incrementBytesRead(val uint64) {
+	d.bytesRead += val
 }
 
 func (di *docValueReader) loadDvChunk(chunkNumber uint64, s *SegmentBase) error {
@@ -289,6 +298,7 @@ func (s *SegmentBase) VisitDocValues(localDocNum uint64, fields []string,
 		if dvs.segment != s {
 			dvs.segment = s
 			dvs.dvrs = nil
+			dvs.bytesRead = 0
 		}
 	}
 
@@ -327,7 +337,9 @@ func (s *SegmentBase) VisitDocValues(localDocNum uint64, fields []string,
 				if err != nil {
 					return dvs, err
 				}
-				s.incrementBytesRead(dvr.BytesRead())
+				dvs.ResetBytesRead(dvr.getBytesRead())
+			} else {
+				dvs.ResetBytesRead(0)
 			}
 
 			_ = dvr.visitDocValues(localDocNum, visitor)

go.mod

@@ -4,9 +4,9 @@ go 1.18
 
 require (
 	github.com/RoaringBitmap/roaring v0.9.4
-	github.com/blevesearch/bleve_index_api v1.0.4
+	github.com/blevesearch/bleve_index_api v1.0.5
 	github.com/blevesearch/mmap-go v1.0.4
-	github.com/blevesearch/scorch_segment_api/v2 v2.1.3
+	github.com/blevesearch/scorch_segment_api/v2 v2.1.4
 	github.com/blevesearch/vellum v1.0.9
 	github.com/golang/snappy v0.0.1
 	github.com/spf13/cobra v1.4.0

go.sum

@@ -2,12 +2,12 @@ github.com/RoaringBitmap/roaring v0.9.4 h1:ckvZSX5gwCRaJYBNe7syNawCU5oruY9gQmjXl
 github.com/RoaringBitmap/roaring v0.9.4/go.mod h1:icnadbWcNyfEHlYdr+tDlOTih1Bf/h+rzPpv4sbomAA=
 github.com/bits-and-blooms/bitset v1.2.0 h1:Kn4yilvwNtMACtf1eYDlG8H77R07mZSPbMjLyS07ChA=
 github.com/bits-and-blooms/bitset v1.2.0/go.mod h1:gIdJ4wp64HaoK2YrL1Q5/N7Y16edYb8uY+O0FJTyyDA=
-github.com/blevesearch/bleve_index_api v1.0.4 h1:mtlzsyJjMIlDngqqB1mq8kPryUMIuEVVbRbJHOWEexU=
-github.com/blevesearch/bleve_index_api v1.0.4/go.mod h1:YXMDwaXFFXwncRS8UobWs7nvo0DmusriM1nztTlj1ms=
+github.com/blevesearch/bleve_index_api v1.0.5 h1:Lc986kpC4Z0/n1g3gg8ul7H+lxgOQPcXb9SxvQGu+tw=
+github.com/blevesearch/bleve_index_api v1.0.5/go.mod h1:YXMDwaXFFXwncRS8UobWs7nvo0DmusriM1nztTlj1ms=
 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.1.3 h1:2UzpR2dR5DvSZk8tVJkcQ7D5xhoK/UBelYw8ttBHrRQ=
-github.com/blevesearch/scorch_segment_api/v2 v2.1.3/go.mod h1:eZrfp1y+lUh+DzFjUcTBUSnKGuunyFIpBIvqYVzJfvc=
+github.com/blevesearch/scorch_segment_api/v2 v2.1.4 h1:LmGmo5twU3gV+natJbKmOktS9eMhokPGKWuR+jX84vk=
+github.com/blevesearch/scorch_segment_api/v2 v2.1.4/go.mod h1:PgVnbbg/t1UkgezPDu8EHLi1BHQ17xUwsFdU6NnOYS0=
 github.com/blevesearch/vellum v1.0.9 h1:PL+NWVk3dDGPCV0hoDu9XLLJgqU4E5s/dOeEJByQ2uQ=
 github.com/blevesearch/vellum v1.0.9/go.mod h1:ul1oT0FhSMDIExNjIxHqJoGpVrBpKCdgDQNxfqgJt7k=
 github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=

intDecoder.go

@@ -17,7 +17,6 @@ package zap
 import (
 	"encoding/binary"
 	"fmt"
-	"sync/atomic"
 )
 
 type chunkedIntDecoder struct {
@@ -59,7 +58,7 @@ func newChunkedIntDecoder(buf []byte, offset uint64, rv *chunkedIntDecoder) *chu
 		rv.chunkOffsets[i], read = binary.Uvarint(buf[offset+n : offset+n+binary.MaxVarintLen64])
 		n += uint64(read)
 	}
-	atomic.AddUint64(&rv.bytesRead, n)
+	rv.bytesRead += n
 	rv.dataStartOffset = offset + n
 	return rv
 }
@@ -70,7 +69,7 @@ func newChunkedIntDecoder(buf []byte, offset uint64, rv *chunkedIntDecoder) *chu
 // the loadChunk retrieves the next chunk and the
 // number of bytes retrieve in that operation is accounted
 func (d *chunkedIntDecoder) getBytesRead() uint64 {
-	return atomic.LoadUint64(&d.bytesRead)
+	return d.bytesRead
 }
 
 func (d *chunkedIntDecoder) loadChunk(chunk int) error {
@@ -89,7 +88,7 @@ func (d *chunkedIntDecoder) loadChunk(chunk int) error {
 	start += s
 	end += e
 	d.curChunkBytes = d.data[start:end]
-	atomic.AddUint64(&d.bytesRead, uint64(len(d.curChunkBytes)))
+	d.bytesRead += uint64(len(d.curChunkBytes))
 	if d.r == nil {
 		d.r = newMemUvarintReader(d.curChunkBytes)
 	} else {
@@ -104,7 +103,7 @@ func (d *chunkedIntDecoder) reset() {
 	d.dataStartOffset = 0
 	d.chunkOffsets = d.chunkOffsets[:0]
 	d.curChunkBytes = d.curChunkBytes[:0]
-	atomic.StoreUint64(&d.bytesRead, 0)
+	d.bytesRead = 0
 	d.data = d.data[:0]
 	if d.r != nil {
 		d.r.Reset([]byte(nil))

posting.go

@@ -19,7 +19,6 @@ import (
 	"fmt"
 	"math"
 	"reflect"
-	"sync/atomic"
 
 	"github.com/RoaringBitmap/roaring"
 	segment "github.com/blevesearch/scorch_segment_api/v2"
@@ -255,15 +254,15 @@ func (p *PostingsList) Count() uint64 {
 // the bytes read from the postings lists stored
 // on disk, while querying
 func (p *PostingsList) ResetBytesRead(val uint64) {
-	atomic.StoreUint64(&p.bytesRead, val)
+	p.bytesRead = val
 }
 
 func (p *PostingsList) BytesRead() uint64 {
-	return atomic.LoadUint64(&p.bytesRead)
+	return p.bytesRead
 }
 
 func (p *PostingsList) incrementBytesRead(val uint64) {
-	atomic.AddUint64(&p.bytesRead, val)
+	p.bytesRead += val
 }
 
 func (p *PostingsList) BytesWritten() uint64 {
@@ -368,15 +367,15 @@ func (i *PostingsIterator) Size() int {
 // the freqNorm and location specific information
 // of a hit
 func (i *PostingsIterator) ResetBytesRead(val uint64) {
-	atomic.StoreUint64(&i.bytesRead, val)
+	i.bytesRead = val
 }
 
 func (i *PostingsIterator) BytesRead() uint64 {
-	return atomic.LoadUint64(&i.bytesRead)
+	return i.bytesRead
 }
 
 func (i *PostingsIterator) incrementBytesRead(val uint64) {
-	atomic.AddUint64(&i.bytesRead, val)
+	i.bytesRead += val
 }
 
 func (i *PostingsIterator) BytesWritten() uint64 {

read.go

@@ -22,8 +22,6 @@ func (s *SegmentBase) getDocStoredMetaAndCompressed(docNum uint64) ([]byte, []by
 	meta := s.mem[storedOffset+n : storedOffset+n+metaLen]
 	data := s.mem[storedOffset+n+metaLen : storedOffset+n+metaLen+dataLen]
 
-	s.incrementBytesRead(metaLen + dataLen)
-
 	return meta, data
 }
 
@@ -41,7 +39,5 @@ func (s *SegmentBase) getDocStoredOffsets(docNum uint64) (
 	dataLen, read := binary.Uvarint(s.mem[storedOffset+n : storedOffset+n+binary.MaxVarintLen64])
 	n += uint64(read)
 
-	s.incrementBytesRead(n)
-
 	return indexOffset, storedOffset, n, metaLen, dataLen
 }

segment.go

@@ -232,8 +232,7 @@ func (s *Segment) ResetBytesRead(val uint64) {
 }
 
 func (s *Segment) BytesRead() uint64 {
-	return atomic.LoadUint64(&s.bytesRead) +
-		atomic.LoadUint64(&s.SegmentBase.bytesRead)
+	return atomic.LoadUint64(&s.bytesRead)
 }
 
 func (s *Segment) BytesWritten() uint64 {
@@ -321,7 +320,7 @@ func (sb *SegmentBase) dictionary(field string) (rv *Dictionary, err error) {
 				// read the length of the vellum data
 				vellumLen, read := binary.Uvarint(sb.mem[dictStart : dictStart+binary.MaxVarintLen64])
 				fstBytes := sb.mem[dictStart+uint64(read) : dictStart+uint64(read)+vellumLen]
-				sb.incrementBytesRead(uint64(read) + vellumLen)
+				rv.incrementBytesRead(uint64(read) + vellumLen)
 				rv.fst, err = vellum.Load(fstBytes)
 				if err != nil {
 					sb.m.Unlock()