github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/committed/skip_prefix_iterator.go (about)

     1  package committed
     2  
     3  import (
     4  	"sort"
     5  	"strings"
     6  
     7  	"github.com/treeverse/lakefs/pkg/graveler"
     8  )
     9  
    10  type rangeValue struct {
    11  	r  *Range
    12  	vr *graveler.ValueRecord
    13  }
    14  
    15  type ImportIterator interface {
    16  	IsCurrentRangeBoundedByPrefix() bool
    17  	IsCurrentPrefixIncludedInRange() bool
    18  }
    19  
    20  type SkipPrefixIterator struct {
    21  	currentPrefixIndex int
    22  	prefixes           []graveler.Prefix
    23  	rangeIterator      Iterator
    24  	currentRangeValue  rangeValue
    25  }
    26  
    27  func (ipi *SkipPrefixIterator) Value() (*graveler.ValueRecord, *Range) {
    28  	return ipi.currentRangeValue.vr, ipi.currentRangeValue.r
    29  }
    30  
    31  func (ipi *SkipPrefixIterator) updateValue() (*graveler.ValueRecord, *Range) {
    32  	vr, r := ipi.rangeIterator.Value()
    33  	ipi.currentRangeValue = rangeValue{
    34  		r,
    35  		vr,
    36  	}
    37  	return vr, r
    38  }
    39  
    40  func (ipi *SkipPrefixIterator) Err() error {
    41  	return ipi.rangeIterator.Err()
    42  }
    43  func (ipi *SkipPrefixIterator) Close() {
    44  	ipi.rangeIterator.Close()
    45  }
    46  func (ipi *SkipPrefixIterator) SeekGE(id graveler.Key) {
    47  	ipi.rangeIterator.SeekGE(id)
    48  }
    49  
    50  func (ipi *SkipPrefixIterator) Next() bool {
    51  	if !ipi.rangeIterator.Next() {
    52  		return false
    53  	}
    54  	vr, r := ipi.updateValue()
    55  	ipi.updatePrefix()
    56  
    57  	if vr == nil && r != nil { // head of range
    58  		for ipi.IsCurrentRangeBoundedByPrefix() {
    59  			if !ipi.rangeIterator.NextRange() {
    60  				return false
    61  			}
    62  			ipi.updateValue()
    63  			ipi.updatePrefix()
    64  		}
    65  	} else {
    66  		prefixLen := len(ipi.prefixes)
    67  		for vr != nil && ipi.currentPrefixIndex < prefixLen && strings.HasPrefix(vr.Key.String(), string(ipi.prefixes[ipi.currentPrefixIndex])) {
    68  			if !ipi.rangeIterator.Next() {
    69  				return false
    70  			}
    71  			vr, _ = ipi.updateValue()
    72  			ipi.updatePrefix()
    73  		}
    74  	}
    75  	return true
    76  }
    77  
    78  func (ipi *SkipPrefixIterator) NextRange() bool {
    79  	if !ipi.rangeIterator.NextRange() {
    80  		return false
    81  	}
    82  	ipi.updateValue()
    83  	ipi.updatePrefix()
    84  
    85  	for ipi.IsCurrentRangeBoundedByPrefix() {
    86  		if !ipi.rangeIterator.NextRange() {
    87  			return false
    88  		}
    89  		ipi.updateValue()
    90  		ipi.updatePrefix()
    91  	}
    92  	return true
    93  }
    94  
    95  func (ipi *SkipPrefixIterator) updatePrefix() {
    96  	if ipi.currentPrefixIndex >= len(ipi.prefixes) {
    97  		return
    98  	}
    99  	currMinKey := string(ipi.currentRangeValue.r.MinKey)
   100  	if ipi.currentRangeValue.vr != nil {
   101  		currMinKey = string(ipi.currentRangeValue.vr.Key)
   102  	}
   103  	// If the current prefix is smaller and isn't the prefix of the currentMinKey, get the next prefix.
   104  	// By the end of this loop, the examined prefix will either be the prefix of the currentMinKey, or
   105  	// lexicographically bigger than it.
   106  	for string(ipi.prefixes[ipi.currentPrefixIndex]) < currMinKey &&
   107  		!strings.HasPrefix(currMinKey, string(ipi.prefixes[ipi.currentPrefixIndex])) {
   108  		p := ipi.currentPrefixIndex + 1
   109  		switch {
   110  		case p >= len(ipi.prefixes): // No more comparable prefixes
   111  			ipi.currentPrefixIndex = p // Block next call for updatePrefix
   112  			return
   113  		case string(ipi.prefixes[p]) <= string(ipi.currentRangeValue.r.MaxKey):
   114  			ipi.currentPrefixIndex = p
   115  		default:
   116  			return
   117  		}
   118  	}
   119  }
   120  
   121  func (ipi *SkipPrefixIterator) getPrefix() *graveler.Prefix {
   122  	if ipi.currentPrefixIndex >= len(ipi.prefixes) {
   123  		return nil
   124  	}
   125  	return &ipi.prefixes[ipi.currentPrefixIndex]
   126  }
   127  
   128  // IsCurrentRangeBoundedByPrefix returns true if both the range's max and min keys have the current prefix.
   129  func (ipi *SkipPrefixIterator) IsCurrentRangeBoundedByPrefix() bool {
   130  	p := ipi.getPrefix()
   131  	if p == nil {
   132  		return false
   133  	}
   134  	r := ipi.currentRangeValue.r
   135  	return strings.HasPrefix(string(r.MinKey), string(*p)) && strings.HasPrefix(string(r.MaxKey), string(*p))
   136  }
   137  
   138  // IsCurrentPrefixIncludedInRange returns true if the examined prefix is either a prefix of the range's min or max key,
   139  // or if the prefix is between the range's min and max keys.
   140  func (ipi *SkipPrefixIterator) IsCurrentPrefixIncludedInRange() bool {
   141  	p := ipi.getPrefix()
   142  	if p == nil {
   143  		return false
   144  	}
   145  	r := ipi.currentRangeValue.r
   146  	inRange := strings.Compare(string(*p), string(r.MinKey)) >= 0 && strings.Compare(string(*p), string(r.MaxKey)) <= 0
   147  	return strings.HasPrefix(string(r.MinKey), string(*p)) || inRange
   148  }
   149  
   150  func NewSkipPrefixIterator(prefixes []graveler.Prefix, rangeIterator Iterator) *SkipPrefixIterator {
   151  	sort.Slice(prefixes, func(i, j int) bool {
   152  		return prefixes[i] < prefixes[j]
   153  	})
   154  	return &SkipPrefixIterator{prefixes: prefixes, currentPrefixIndex: 0, rangeIterator: rangeIterator}
   155  }