github.com/weaviate/sroar@v0.0.0-20230210105426-26108af5465d/iterator.go (about)

     1  /*
     2   * Copyright 2021 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package sroar
    18  
    19  import (
    20  	"math/bits"
    21  )
    22  
    23  type Iterator struct {
    24  	bm *Bitmap
    25  
    26  	keys   []uint64
    27  	keyIdx int
    28  
    29  	contIdx int
    30  
    31  	bitmapIdx int
    32  	bitset    uint16
    33  }
    34  
    35  func (bm *Bitmap) NewRangeIterators(numRanges int) []*Iterator {
    36  	keyn := bm.keys.numKeys()
    37  	iters := make([]*Iterator, numRanges)
    38  	width := keyn / numRanges
    39  	rem := keyn % numRanges
    40  	cnt := 0
    41  
    42  	// This loop distributes the key equally to the ranges. For example: If numRanges = 3
    43  	// and keyn = 8 then it will be distributes as [3, 3, 2]
    44  	for i := 0; i < numRanges; i++ {
    45  		iters[i] = bm.NewIterator()
    46  		n := width
    47  		if i < rem {
    48  			n = width + 1
    49  		}
    50  		iters[i].keys = iters[i].keys[cnt : cnt+2*n]
    51  		cnt = cnt + 2*n
    52  	}
    53  	return iters
    54  }
    55  
    56  func (bm *Bitmap) NewIterator() *Iterator {
    57  	return &Iterator{
    58  		bm:        bm,
    59  		keys:      bm.keys[indexNodeStart : indexNodeStart+bm.keys.numKeys()*2],
    60  		keyIdx:    0,
    61  		contIdx:   -1,
    62  		bitmapIdx: -1,
    63  	}
    64  }
    65  
    66  func (it *Iterator) Next() uint64 {
    67  	if len(it.keys) == 0 {
    68  		return 0
    69  	}
    70  
    71  	key := it.keys[it.keyIdx]
    72  	off := it.keys[it.keyIdx+1]
    73  	cont := it.bm.getContainer(off)
    74  	card := getCardinality(cont)
    75  
    76  	// Loop until we find a container on which next operation is possible. When such a container
    77  	// is found, reset the variables responsible for container iteration.
    78  	for card == 0 || it.contIdx+1 >= card {
    79  		if it.keyIdx+2 >= len(it.keys) {
    80  			return 0
    81  		}
    82  		// jump by 2 because key is followed by a value
    83  		it.keyIdx += 2
    84  		it.contIdx = -1
    85  		it.bitmapIdx = -1
    86  		it.bitset = 0
    87  		key = it.keys[it.keyIdx]
    88  		off = it.keys[it.keyIdx+1]
    89  		cont = it.bm.getContainer(off)
    90  		card = getCardinality(cont)
    91  	}
    92  
    93  	//  The above loop assures that we can do next in this container.
    94  	it.contIdx++
    95  	switch cont[indexType] {
    96  	case typeArray:
    97  		return key | uint64(cont[int(startIdx)+it.contIdx])
    98  	case typeBitmap:
    99  		// A bitmap container is an array of uint16s.
   100  		// If the container is bitmap, go to the index which has a non-zero value.
   101  		for it.bitset == 0 && it.bitmapIdx+1 < len(cont[startIdx:]) {
   102  			it.bitmapIdx++
   103  			it.bitset = cont[int(startIdx)+it.bitmapIdx]
   104  		}
   105  		assert(it.bitset > 0)
   106  
   107  		// msbIdx is the index of most-significant bit. In this iteration we choose this set bit
   108  		// and make it zero.
   109  		msbIdx := uint16(bits.LeadingZeros16(it.bitset))
   110  		msb := 1 << (16 - msbIdx - 1)
   111  		it.bitset ^= uint16(msb)
   112  		return key | uint64(it.bitmapIdx*16+int(msbIdx))
   113  	}
   114  	return 0
   115  }
   116  
   117  type ManyItr struct {
   118  	index int
   119  	arr   []uint64
   120  }
   121  
   122  // TODO: See if this is needed, we should remove this
   123  func (r *Bitmap) ManyIterator() *ManyItr {
   124  	return &ManyItr{
   125  		arr: r.ToArray(),
   126  	}
   127  
   128  }
   129  
   130  func (itr *ManyItr) NextMany(buf []uint64) int {
   131  	count := 0
   132  	for i := 0; i < len(buf); i++ {
   133  		if itr.index == len(itr.arr) {
   134  			break
   135  		}
   136  		buf[i] = itr.arr[itr.index]
   137  		itr.index++
   138  		count++
   139  	}
   140  	return count
   141  }