github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/storage/index/aggregate_iter.go (about)

     1  // Copyright (c) 2021 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package index
    22  
    23  import (
    24  	"github.com/m3db/m3/src/m3ninx/index/segment"
    25  	"github.com/m3db/m3/src/x/context"
    26  	xerrors "github.com/m3db/m3/src/x/errors"
    27  )
    28  
    29  var _ AggregateIterator = &aggregateIter{}
    30  
    31  type aggregateIter struct {
    32  	// immutable state
    33  	readers     []segment.Reader
    34  	iterateOpts fieldsAndTermsIteratorOpts
    35  	newIterFn   newFieldsAndTermsIteratorFn
    36  
    37  	// immutable state after first Next() call
    38  	iters []fieldsAndTermsIterator
    39  
    40  	// mutable state
    41  	idx                    int
    42  	err                    error
    43  	done                   bool
    44  	currField, currTerm    []byte
    45  	nextField, nextTerm    []byte
    46  	docsCount, seriesCount int
    47  }
    48  
    49  func (it *aggregateIter) Next(ctx context.Context) bool {
    50  	if it.Done() {
    51  		return false
    52  	}
    53  	if it.iters == nil {
    54  		for _, reader := range it.readers {
    55  			iter, err := it.newIterFn(ctx, reader, it.iterateOpts)
    56  			if err != nil {
    57  				it.err = err
    58  				return false
    59  			}
    60  			it.iters = append(it.iters, iter)
    61  		}
    62  		if !it.next() {
    63  			it.done = true
    64  			return false
    65  		}
    66  		it.nextField, it.nextTerm = it.current()
    67  	}
    68  	// the fieldAndTermsIterator mutates the underlying byte slice, so we need to copy to preserve the value.
    69  	it.currField = append(it.currField[:0], it.nextField...)
    70  	it.currTerm = append(it.currTerm[:0], it.nextTerm...)
    71  
    72  	if it.next() {
    73  		it.nextField, it.nextTerm = it.current()
    74  	} else {
    75  		// the iterators have been exhausted. mark done so the next call Done returns true. Still return true from
    76  		// this call so the caller can retrieve the last element with Current.
    77  		it.done = true
    78  	}
    79  	return true
    80  }
    81  
    82  func (it *aggregateIter) next() bool {
    83  	if it.idx == len(it.iters) {
    84  		return false
    85  	}
    86  	currIter := it.iters[it.idx]
    87  
    88  	for !currIter.Next() {
    89  		if err := currIter.Err(); err != nil {
    90  			it.err = err
    91  			return false
    92  		}
    93  		// move to next iterator so we don't try to Close twice.
    94  		it.idx++
    95  
    96  		if err := currIter.Close(); err != nil {
    97  			it.err = err
    98  			return false
    99  		}
   100  
   101  		if it.idx == len(it.iters) {
   102  			return false
   103  		}
   104  		currIter = it.iters[it.idx]
   105  	}
   106  	return true
   107  }
   108  
   109  func (it *aggregateIter) current() (field, term []byte) {
   110  	return it.iters[it.idx].Current()
   111  }
   112  
   113  func (it *aggregateIter) Err() error {
   114  	return it.err
   115  }
   116  
   117  func (it *aggregateIter) Done() bool {
   118  	return it.err != nil || it.done
   119  }
   120  
   121  func (it *aggregateIter) Current() (field, term []byte) {
   122  	return it.currField, it.currTerm
   123  }
   124  
   125  func (it *aggregateIter) fieldsAndTermsIteratorOpts() fieldsAndTermsIteratorOpts {
   126  	return it.iterateOpts
   127  }
   128  
   129  func (it *aggregateIter) Close() error {
   130  	if it.iters == nil {
   131  		return nil
   132  	}
   133  	var multiErr xerrors.MultiError
   134  	// close any iterators that weren't closed in Next.
   135  	for i := it.idx; i < len(it.iters); i++ {
   136  		multiErr = multiErr.Add(it.iters[i].Close())
   137  	}
   138  	return multiErr.FinalError()
   139  }
   140  
   141  func (it *aggregateIter) AddSeries(count int) {
   142  	it.seriesCount += count
   143  }
   144  
   145  func (it *aggregateIter) AddDocs(count int) {
   146  	it.docsCount += count
   147  }
   148  
   149  func (it *aggregateIter) Counts() (series, docs int) {
   150  	return it.seriesCount, it.docsCount
   151  }