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 }