github.com/weaviate/sroar@v0.0.0-20230210105426-26108af5465d/bitmap.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 "fmt" 21 "math" 22 "sort" 23 "strings" 24 "sync" 25 26 "github.com/pkg/errors" 27 ) 28 29 var empty = make([]uint16, 16<<20) 30 31 const mask = uint64(0xFFFFFFFFFFFF0000) 32 33 type Bitmap struct { 34 data []uint16 35 keys node 36 37 // This _ptr is only used when we start with a []byte instead of a 38 // []uint16. Because we do an unsafe conversion to []uint16 data, and hence, 39 // do NOT own a valid pointer to the underlying array. 40 _ptr []byte 41 42 // memMoved keeps track of how many uint16 moves we had to do. The smaller 43 // this number, the more efficient we have been. 44 memMoved int 45 } 46 47 // FromBuffer returns a pointer to bitmap corresponding to the given buffer. This bitmap shouldn't 48 // be modified because it might corrupt the given buffer. 49 func FromBuffer(data []byte) *Bitmap { 50 assert(len(data)%2 == 0) 51 if len(data) < 8 { 52 return NewBitmap() 53 } 54 du := toUint16Slice(data) 55 x := toUint64Slice(du[:4])[indexNodeSize] 56 return &Bitmap{ 57 data: du, 58 _ptr: data, // Keep a hold of data, otherwise GC would do its thing. 59 keys: toUint64Slice(du[:x]), 60 } 61 } 62 63 // FromBufferWithCopy creates a copy of the given buffer and returns a bitmap based on the copied 64 // buffer. This bitmap is safe for both read and write operations. 65 func FromBufferWithCopy(src []byte) *Bitmap { 66 assert(len(src)%2 == 0) 67 if len(src) < 8 { 68 return NewBitmap() 69 } 70 src16 := toUint16Slice(src) 71 dst16 := make([]uint16, len(src16)) 72 copy(dst16, src16) 73 x := toUint64Slice(dst16[:4])[indexNodeSize] 74 75 return &Bitmap{ 76 data: dst16, 77 keys: toUint64Slice(dst16[:x]), 78 } 79 } 80 81 func (ra *Bitmap) ToBuffer() []byte { 82 if ra.IsEmpty() { 83 return nil 84 } 85 return toByteSlice(ra.data) 86 } 87 88 func (ra *Bitmap) ToBufferWithCopy() []byte { 89 if ra.IsEmpty() { 90 return nil 91 } 92 buf := make([]uint16, len(ra.data)) 93 copy(buf, ra.data) 94 return toByteSlice(buf) 95 } 96 97 func NewBitmap() *Bitmap { 98 return NewBitmapWith(2) 99 } 100 101 func NewBitmapWith(numKeys int) *Bitmap { 102 if numKeys < 2 { 103 panic("Must contain at least two keys.") 104 } 105 ra := &Bitmap{ 106 // Each key must also keep an offset. So, we need to double the number 107 // of uint64s allocated. Plus, we need to make space for the first 2 108 // uint64s to store the number of keys and node size. 109 data: make([]uint16, 4*(2*numKeys+2)), 110 } 111 ra.keys = toUint64Slice(ra.data) 112 ra.keys.setNodeSize(len(ra.data)) 113 114 // Always generate a container for key = 0x00. Otherwise, node gets confused 115 // about whether a zero key is a new key or not. 116 offset := ra.newContainer(minContainerSize) 117 // First two are for num keys. index=2 -> 0 key. index=3 -> offset. 118 ra.keys.setAt(indexNodeStart+1, offset) 119 ra.keys.setNumKeys(1) 120 121 return ra 122 } 123 124 func (ra *Bitmap) initSpaceForKeys(N int) { 125 if N == 0 { 126 return 127 } 128 curSize := uint64(len(ra.keys) * 4) // U64 -> U16 129 bySize := uint64(N * 8) // 2xU64 (key, value) -> 2x4xU16 130 131 // The following code is borrowed from setKey. 132 ra.scootRight(curSize, bySize) 133 ra.keys = toUint64Slice(ra.data[:curSize+bySize]) 134 ra.keys.setNodeSize(int(curSize + bySize)) 135 assert(1 == ra.keys.numKeys()) // This initialization assumes that the number of keys are 1. 136 137 // The containers have moved to the right bySize. So, update their offsets. 138 // Currently, there's only one container. 139 val := ra.keys.val(0) 140 ra.keys.setAt(valOffset(0), val+uint64(bySize)) 141 } 142 143 // setKey sets a key and container offset. 144 func (ra *Bitmap) setKey(k uint64, offset uint64) uint64 { 145 if added := ra.keys.set(k, offset); !added { 146 // No new key was added. So, we can just return. 147 return offset 148 } 149 // A new key was added. Let's ensure that ra.keys is not full. 150 if !ra.keys.isFull() { 151 return offset 152 } 153 154 // ra.keys is full. We should expand its size. 155 curSize := uint64(len(ra.keys) * 4) // Multiply by 4 for U64 -> U16. 156 bySize := curSize 157 if bySize > math.MaxUint16 { 158 bySize = math.MaxUint16 159 } 160 161 ra.scootRight(curSize, bySize) 162 ra.keys = toUint64Slice(ra.data[:curSize+bySize]) 163 ra.keys.setNodeSize(int(curSize + bySize)) 164 165 // All containers have moved to the right by bySize bytes. 166 // Update their offsets. 167 n := ra.keys 168 for i := 0; i < n.maxKeys(); i++ { 169 val := n.val(i) 170 if val > 0 { 171 n.setAt(valOffset(i), val+uint64(bySize)) 172 } 173 } 174 return offset + bySize 175 } 176 177 func (ra *Bitmap) fastExpand(bySize uint64) { 178 prev := len(ra.keys) * 4 // Multiply by 4 to convert from u16 to u64. 179 180 // This following statement also works. But, given how much fastExpand gets 181 // called (a lot), probably better to control allocation. 182 // ra.data = append(ra.data, empty[:bySize]...) 183 184 toSize := len(ra.data) + int(bySize) 185 if toSize <= cap(ra.data) { 186 ra.data = ra.data[:toSize] 187 return 188 } 189 growBy := cap(ra.data) 190 if growBy < int(bySize) { 191 growBy = int(bySize) 192 } 193 out := make([]uint16, cap(ra.data)+growBy) 194 copy(out, ra.data) 195 ra.data = out[:toSize] 196 ra._ptr = nil // Allow Go to GC whatever this was pointing to. 197 // Re-reference ra.keys correctly because underlying array has changed. 198 ra.keys = toUint64Slice(ra.data[:prev]) 199 } 200 201 // scootRight isn't aware of containers. It's going to create empty space of 202 // bySize at the given offset in ra.data. The offset doesn't need to line up 203 // with a container. 204 func (ra *Bitmap) scootRight(offset uint64, bySize uint64) { 205 left := ra.data[offset:] 206 207 ra.fastExpand(bySize) // Expand the buffer. 208 right := ra.data[len(ra.data)-len(left):] 209 n := copy(right, left) // Move data right. 210 ra.memMoved += n 211 212 Memclr(ra.data[offset : offset+uint64(bySize)]) // Zero out the space in the middle. 213 } 214 215 // scootLeft removes size number of uint16s starting from the given offset. 216 func (ra *Bitmap) scootLeft(offset uint64, size uint64) { 217 n := uint64(len(ra.data)) 218 right := ra.data[offset+size:] 219 ra.memMoved += copy(ra.data[offset:], right) 220 ra.data = ra.data[:n-size] 221 } 222 223 func (ra *Bitmap) newContainer(sz uint16) uint64 { 224 offset := uint64(len(ra.data)) 225 ra.fastExpand(uint64(sz)) 226 Memclr(ra.data[offset : offset+uint64(sz)]) 227 ra.data[offset] = sz 228 return offset 229 } 230 231 // expandContainer would expand a container at the given offset. It would typically double the size 232 // of the container, until it reaches a threshold, where the size of the container would reach 2^16. 233 // Expressed in uint16s, that'd be (2^16)/(2^4) = 2^12 = 4096. So, if the container size >= 2048, 234 // then doubling that would put it above 4096. That's why in the code below, you see the checks for 235 // size 2048. 236 func (ra *Bitmap) expandContainer(offset uint64) { 237 sz := ra.data[offset] 238 if sz == 0 { 239 panic("Container size should NOT be zero") 240 } 241 bySize := uint16(sz) 242 if sz >= 2048 { 243 // Size is in uint16. Half of max allowed size. If we're expanding the container by more 244 // than 2048, we should just cap it to max size of 4096. 245 assert(sz < maxContainerSize) 246 bySize = maxContainerSize - sz 247 } 248 249 // Select the portion to the right of the container, beyond its right boundary. 250 ra.scootRight(offset+uint64(sz), uint64(bySize)) 251 ra.keys.updateOffsets(offset, uint64(bySize), true) 252 253 if sz < 2048 { 254 ra.data[offset] = sz + bySize 255 256 } else { 257 // Convert to bitmap container. 258 src := array(ra.getContainer(offset)) 259 buf := src.toBitmapContainer(nil) 260 assert(copy(ra.data[offset:], buf) == maxContainerSize) 261 } 262 } 263 264 // stepSize is used for container expansion. For a container of given size n, 265 // stepSize would return the target size. This function is used to reduce the 266 // number of times expansion needs to happen for each container. 267 func stepSize(n uint16) uint16 { 268 // <=64 -> 128 269 // <=128 -> 256 270 // <=256 -> 512 271 // <=512 -> 1024 272 // <=1024 -> 2048 273 // >1024 -> maxSize (convert to bitmap) 274 for i := uint16(64); i <= 1024; i *= 2 { 275 if n <= i { 276 return i * 2 277 } 278 } 279 return maxContainerSize 280 } 281 282 // copyAt would copy over a given container via src, into the container at 283 // offset. If src is a bitmap, it would copy it over directly. If src is an 284 // array container, then it would follow these paths: 285 // - If src is smaller than dst, copy it over. 286 // - If not, look for target size for dst using the stepSize function. 287 // - If target size is maxSize, then convert src to a bitmap container, and 288 // copy to dst. 289 // - If target size is not max size, then expand dst container and copy src. 290 func (ra *Bitmap) copyAt(offset uint64, src []uint16) { 291 dstSize := ra.data[offset] 292 if dstSize == 0 { 293 panic("Container size should NOT be zero") 294 } 295 296 // The src is a bitmapContainer. Just copy it over. 297 if src[indexType] == typeBitmap { 298 assert(src[indexSize] == maxContainerSize) 299 bySize := uint16(maxContainerSize) - dstSize 300 // Select the portion to the right of the container, beyond its right boundary. 301 ra.scootRight(offset+uint64(dstSize), uint64(bySize)) 302 ra.keys.updateOffsets(offset, uint64(bySize), true) 303 assert(copy(ra.data[offset:], src) == len(src)) 304 return 305 } 306 307 // src is an array container. Check if dstSize >= src. If so, just copy. 308 // But, do keep dstSize intact, otherwise we'd lose portion of our container. 309 if dstSize >= src[indexSize] { 310 assert(copy(ra.data[offset:], src) == len(src)) 311 ra.data[offset] = dstSize 312 return 313 } 314 315 // dstSize < src. Determine the target size of the container. 316 targetSz := stepSize(dstSize) 317 for targetSz < src[indexSize] { 318 targetSz = stepSize(targetSz) 319 } 320 321 if targetSz == maxContainerSize { 322 // Looks like the targetSize is now maxSize. So, convert src to bitmap container. 323 s := array(src) 324 325 bySize := uint16(maxContainerSize) - dstSize 326 // Select the portion to the right of the container, beyond its right boundary. 327 ra.scootRight(offset+uint64(dstSize), uint64(bySize)) 328 ra.keys.updateOffsets(offset, uint64(bySize), true) 329 330 // Update the space of the container, so getContainer would work correctly. 331 ra.data[offset] = maxContainerSize 332 333 // Convert the src array to bitmap and write it directly over to the container. 334 out := ra.getContainer(offset) 335 Memclr(out) 336 s.toBitmapContainer(out) 337 return 338 } 339 340 // targetSize is not maxSize. Let's expand to targetSize and copy array. 341 bySize := targetSz - dstSize 342 ra.scootRight(offset+uint64(dstSize), uint64(bySize)) 343 ra.keys.updateOffsets(offset, uint64(bySize), true) 344 assert(copy(ra.data[offset:], src) == len(src)) 345 ra.data[offset] = targetSz 346 } 347 348 func (ra Bitmap) getContainer(offset uint64) []uint16 { 349 data := ra.data[offset:] 350 if len(data) == 0 { 351 panic(fmt.Sprintf("No container found at offset: %d\n", offset)) 352 } 353 sz := data[0] 354 return data[:sz] 355 } 356 357 func (ra *Bitmap) Clone() *Bitmap { 358 abuf := ra.ToBuffer() 359 bbuf := make([]byte, len(abuf)) 360 copy(bbuf, abuf) 361 return FromBuffer(bbuf) 362 } 363 364 func (ra *Bitmap) IsEmpty() bool { 365 if ra == nil { 366 return true 367 } 368 N := ra.keys.numKeys() 369 for i := 0; i < N; i++ { 370 offset := ra.keys.val(i) 371 cont := ra.getContainer(offset) 372 if c := getCardinality(cont); c > 0 { 373 return false 374 } 375 } 376 return true 377 } 378 379 func (ra *Bitmap) Set(x uint64) bool { 380 key := x & mask 381 offset, has := ra.keys.getValue(key) 382 if !has { 383 // We need to add a container. 384 o := ra.newContainer(minContainerSize) 385 // offset might have been updated by setKey. 386 offset = ra.setKey(key, o) 387 } 388 c := ra.getContainer(offset) 389 switch c[indexType] { 390 case typeArray: 391 p := array(c) 392 if added := p.add(uint16(x)); !added { 393 return false 394 } 395 if p.isFull() { 396 ra.expandContainer(offset) 397 } 398 return true 399 case typeBitmap: 400 b := bitmap(c) 401 return b.add(uint16(x)) 402 } 403 panic("we shouldn't reach here") 404 } 405 406 func FromSortedList(vals []uint64) *Bitmap { 407 var arr []uint16 408 var hi, lastHi, off uint64 409 410 ra := NewBitmap() 411 412 if len(vals) == 0 { 413 return ra 414 } 415 416 // Set the keys beforehand so that we don't need to move a lot of memory because of adding keys. 417 var numKeys int 418 for _, x := range vals { 419 hi = x & mask 420 if hi != 0 && hi != lastHi { 421 numKeys++ 422 } 423 lastHi = hi 424 } 425 ra.initSpaceForKeys(numKeys) 426 427 finalize := func(l []uint16, key uint64) { 428 if len(l) == 0 { 429 return 430 } 431 if len(l) <= 2048 { 432 // 4 uint16s for the header, and extra 4 uint16s so that adding more elements using 433 // Set operation doesn't fail. 434 sz := uint16(8 + len(l)) 435 off = ra.newContainer(sz) 436 c := ra.getContainer(off) 437 c[indexSize] = sz 438 c[indexType] = typeArray 439 setCardinality(c, len(l)) 440 for i := 0; i < len(l); i++ { 441 c[int(startIdx)+i] = l[i] 442 } 443 444 } else { 445 off = ra.newContainer(maxContainerSize) 446 c := ra.getContainer(off) 447 c[indexSize] = maxContainerSize 448 c[indexType] = typeBitmap 449 for _, v := range l { 450 bitmap(c).add(v) 451 } 452 } 453 ra.setKey(key, off) 454 return 455 } 456 457 lastHi = 0 458 for _, x := range vals { 459 hi = x & mask 460 // Finalize the last container before proceeding ahead 461 if hi != 0 && hi != lastHi { 462 finalize(arr, lastHi) 463 arr = arr[:0] 464 } 465 arr = append(arr, uint16(x)) 466 lastHi = hi 467 } 468 finalize(arr, lastHi) 469 return ra 470 } 471 472 // TODO: Potentially this can be optimized. 473 func (ra *Bitmap) SetMany(vals []uint64) { 474 for _, k := range vals { 475 ra.Set(k) 476 } 477 } 478 479 // Select returns the element at the xth index. (0-indexed) 480 func (ra *Bitmap) Select(x uint64) (uint64, error) { 481 if x >= uint64(ra.GetCardinality()) { 482 return 0, errors.Errorf("index %d is not less than the cardinality: %d", 483 x, ra.GetCardinality()) 484 } 485 n := ra.keys.numKeys() 486 for i := 0; i < n; i++ { 487 off := ra.keys.val(i) 488 con := ra.getContainer(off) 489 c := uint64(getCardinality(con)) 490 assert(c != uint64(invalidCardinality)) 491 if x < c { 492 key := ra.keys.key(i) 493 switch con[indexType] { 494 case typeArray: 495 return key | uint64(array(con).all()[x]), nil 496 case typeBitmap: 497 return key | uint64(bitmap(con).selectAt(int(x))), nil 498 } 499 } 500 x -= c 501 } 502 panic("should not reach here") 503 } 504 505 func (ra *Bitmap) Contains(x uint64) bool { 506 if ra == nil { 507 return false 508 } 509 key := x & mask 510 offset, has := ra.keys.getValue(key) 511 if !has { 512 return false 513 } 514 y := uint16(x) 515 516 c := ra.getContainer(offset) 517 switch c[indexType] { 518 case typeArray: 519 p := array(c) 520 return p.has(y) 521 case typeBitmap: 522 b := bitmap(c) 523 return b.has(y) 524 } 525 return false 526 } 527 528 func (ra *Bitmap) Remove(x uint64) bool { 529 if ra == nil { 530 return false 531 } 532 key := x & mask 533 offset, has := ra.keys.getValue(key) 534 if !has { 535 return false 536 } 537 c := ra.getContainer(offset) 538 switch c[indexType] { 539 case typeArray: 540 p := array(c) 541 return p.remove(uint16(x)) 542 case typeBitmap: 543 b := bitmap(c) 544 return b.remove(uint16(x)) 545 } 546 return true 547 } 548 549 // Remove range removes [lo, hi) from the bitmap. 550 func (ra *Bitmap) RemoveRange(lo, hi uint64) { 551 if lo > hi { 552 panic("lo should not be more than hi") 553 } 554 if lo == hi { 555 return 556 } 557 558 k1 := lo & mask 559 k2 := hi & mask 560 561 defer ra.Cleanup() 562 563 // Complete range lie in a single container 564 if k1 == k2 { 565 if off, has := ra.keys.getValue(k1); has { 566 c := ra.getContainer(off) 567 removeRangeContainer(c, uint16(lo), uint16(hi)-1) 568 } 569 return 570 } 571 572 // Remove all the containers in range [k1+1, k2-1]. 573 n := ra.keys.numKeys() 574 st := ra.keys.search(k1) 575 key := ra.keys.key(st) 576 if key == k1 { 577 st++ 578 } 579 580 for i := st; i < n; i++ { 581 key := ra.keys.key(i) 582 if key >= k2 { 583 break 584 } 585 if off, has := ra.keys.getValue(key); has { 586 zeroOutContainer(ra.getContainer(off)) 587 } 588 } 589 590 // Remove elements >= lo in k1's container 591 if off, has := ra.keys.getValue(k1); has { 592 c := ra.getContainer(off) 593 if uint16(lo) == 0 { 594 zeroOutContainer(c) 595 } else { 596 removeRangeContainer(c, uint16(lo), math.MaxUint16) 597 } 598 } 599 600 if uint16(hi) == 0 { 601 return 602 } 603 604 // Remove all elements < hi in k2's container 605 if off, has := ra.keys.getValue(k2); has { 606 c := ra.getContainer(off) 607 removeRangeContainer(c, 0, uint16(hi)-1) 608 } 609 } 610 611 func (ra *Bitmap) Reset() { 612 // reset ra.data to size enough for one container and corresponding key. 613 // 2 u64 is needed for header and another 2 u16 for the key 0. 614 ra.data = ra.data[:16+minContainerSize] 615 ra.keys = toUint64Slice(ra.data) 616 617 offset := ra.newContainer(minContainerSize) 618 ra.keys.setAt(indexNodeStart+1, offset) 619 ra.keys.setNumKeys(1) 620 } 621 622 func (ra *Bitmap) GetCardinality() int { 623 if ra == nil { 624 return 0 625 } 626 N := ra.keys.numKeys() 627 var sz int 628 for i := 0; i < N; i++ { 629 offset := ra.keys.val(i) 630 c := ra.getContainer(offset) 631 sz += getCardinality(c) 632 } 633 return sz 634 } 635 636 func (ra *Bitmap) ToArray() []uint64 { 637 if ra == nil { 638 return nil 639 } 640 res := make([]uint64, 0, ra.GetCardinality()) 641 N := ra.keys.numKeys() 642 for i := 0; i < N; i++ { 643 key := ra.keys.key(i) 644 off := ra.keys.val(i) 645 c := ra.getContainer(off) 646 647 switch c[indexType] { 648 case typeArray: 649 a := array(c) 650 for _, lo := range a.all() { 651 res = append(res, key|uint64(lo)) 652 } 653 case typeBitmap: 654 b := bitmap(c) 655 out := b.all() 656 for _, x := range out { 657 res = append(res, key|uint64(x)) 658 } 659 } 660 } 661 return res 662 } 663 664 func (ra *Bitmap) String() string { 665 var b strings.Builder 666 b.WriteRune('\n') 667 668 var usedSize, card int 669 usedSize += 4 * (ra.keys.numKeys()) 670 for i := 0; i < ra.keys.numKeys(); i++ { 671 k := ra.keys.key(i) 672 v := ra.keys.val(i) 673 c := ra.getContainer(v) 674 675 sz := c[indexSize] 676 usedSize += int(sz) 677 card += getCardinality(c) 678 679 b.WriteString(fmt.Sprintf( 680 "[%03d] Key: %#8x. Offset: %7d. Size: %4d. Type: %d. Card: %6d. Uint16/Uid: %.2f\n", 681 i, k, v, sz, c[indexType], getCardinality(c), float64(sz)/float64(getCardinality(c)))) 682 } 683 b.WriteString(fmt.Sprintf("Number of containers: %d. Cardinality: %d\n", 684 ra.keys.numKeys(), card)) 685 686 amp := float64(len(ra.data)-usedSize) / float64(usedSize) 687 b.WriteString(fmt.Sprintf( 688 "Size in Uint16s. Used: %d. Total: %d. Space Amplification: %.2f%%. Moved: %.2fx\n", 689 usedSize, len(ra.data), amp*100.0, float64(ra.memMoved)/float64(usedSize))) 690 691 b.WriteString(fmt.Sprintf("Used Uint16/Uid: %.2f. Total Uint16/Uid: %.2f", 692 float64(usedSize)/float64(card), float64(len(ra.data))/float64(card))) 693 694 return b.String() 695 } 696 697 const fwd int = 0x01 698 const rev int = 0x02 699 700 func (ra *Bitmap) Minimum() uint64 { return ra.extreme(fwd) } 701 func (ra *Bitmap) Maximum() uint64 { return ra.extreme(rev) } 702 703 func (ra *Bitmap) Debug(x uint64) string { 704 var b strings.Builder 705 hi := x & mask 706 off, found := ra.keys.getValue(hi) 707 if !found { 708 b.WriteString(fmt.Sprintf("Unable to find the container for x: %#x\n", hi)) 709 b.WriteString(ra.String()) 710 } 711 c := ra.getContainer(off) 712 lo := uint16(x) 713 714 b.WriteString(fmt.Sprintf("x: %#x lo: %#x. offset: %d\n", x, lo, off)) 715 716 switch c[indexType] { 717 case typeArray: 718 case typeBitmap: 719 idx := lo / 16 720 pos := lo % 16 721 b.WriteString(fmt.Sprintf("At idx: %d. Pos: %d val: %#b\n", idx, pos, c[startIdx+idx])) 722 } 723 return b.String() 724 } 725 726 func (ra *Bitmap) extreme(dir int) uint64 { 727 N := ra.keys.numKeys() 728 if N == 0 { 729 return 0 730 } 731 732 var k uint64 733 var c []uint16 734 735 if dir == fwd { 736 for i := 0; i < N; i++ { 737 offset := ra.keys.val(i) 738 c = ra.getContainer(offset) 739 if getCardinality(c) > 0 { 740 k = ra.keys.key(i) 741 break 742 } 743 } 744 } else { 745 for i := N - 1; i >= 0; i-- { 746 offset := ra.keys.val(i) 747 c = ra.getContainer(offset) 748 if getCardinality(c) > 0 { 749 k = ra.keys.key(i) 750 break 751 } 752 } 753 } 754 755 switch c[indexType] { 756 case typeArray: 757 a := array(c) 758 if dir == fwd { 759 return k | uint64(a.minimum()) 760 } 761 return k | uint64(a.maximum()) 762 case typeBitmap: 763 b := bitmap(c) 764 if dir == fwd { 765 return k | uint64(b.minimum()) 766 } 767 return k | uint64(b.maximum()) 768 default: 769 panic("We don't support this type of container") 770 } 771 } 772 773 func (ra *Bitmap) And(bm *Bitmap) { 774 if bm == nil { 775 ra.Reset() 776 return 777 } 778 779 a, b := ra, bm 780 ai, an := 0, a.keys.numKeys() 781 bi, bn := 0, b.keys.numKeys() 782 783 for ai < an && bi < bn { 784 ak := a.keys.key(ai) 785 bk := b.keys.key(bi) 786 if ak == bk { 787 off := a.keys.val(ai) 788 ac := a.getContainer(off) 789 790 off = b.keys.val(bi) 791 bc := b.getContainer(off) 792 793 // do the intersection 794 // TODO: See if we can do containerAnd operation in-place. 795 c := containerAnd(ac, bc) 796 797 // create a new container and update the key offset to this container. 798 offset := a.newContainer(uint16(len(c))) 799 copy(a.data[offset:], c) 800 a.setKey(ak, offset) 801 ai++ 802 bi++ 803 } else if ak < bk { 804 off := a.keys.val(ai) 805 zeroOutContainer(a.getContainer(off)) 806 ai++ 807 } else { 808 bi++ 809 } 810 } 811 for ai < an { 812 off := a.keys.val(ai) 813 zeroOutContainer(a.getContainer(off)) 814 ai++ 815 } 816 } 817 818 func And(a, b *Bitmap) *Bitmap { 819 ai, an := 0, a.keys.numKeys() 820 bi, bn := 0, b.keys.numKeys() 821 822 res := NewBitmap() 823 for ai < an && bi < bn { 824 ak := a.keys.key(ai) 825 bk := a.keys.key(bi) 826 if ak == bk { 827 // Do the intersection. 828 off := a.keys.val(ai) 829 ac := a.getContainer(off) 830 831 off = b.keys.val(bi) 832 bc := b.getContainer(off) 833 834 outc := containerAnd(ac, bc) 835 if getCardinality(outc) > 0 { 836 offset := res.newContainer(uint16(len(outc))) 837 copy(res.data[offset:], outc) 838 res.setKey(ak, offset) 839 } 840 ai++ 841 bi++ 842 } else if ak < bk { 843 ai++ 844 } else { 845 bi++ 846 } 847 } 848 return res 849 } 850 851 func (ra *Bitmap) AndNot(bm *Bitmap) { 852 if bm == nil { 853 return 854 } 855 a, b := ra, bm 856 var ai, bi int 857 858 buf := make([]uint16, maxContainerSize) 859 for ai < a.keys.numKeys() && bi < b.keys.numKeys() { 860 ak := a.keys.key(ai) 861 bk := b.keys.key(bi) 862 if ak == bk { 863 off := a.keys.val(ai) 864 ac := a.getContainer(off) 865 866 off = b.keys.val(bi) 867 bc := b.getContainer(off) 868 869 // TODO: See if we can do containerAndNot operation in-place. 870 c := containerAndNot(ac, bc, buf) 871 // create a new container and update the key offset to this container. 872 offset := a.newContainer(uint16(len(c))) 873 copy(a.data[offset:], c) 874 a.setKey(ak, offset) 875 876 ai++ 877 bi++ 878 continue 879 } 880 if ak < bk { 881 ai++ 882 } else { 883 bi++ 884 } 885 } 886 } 887 888 // TODO: Check if we want to use lazyMode 889 func (dst *Bitmap) Or(src *Bitmap) { 890 if src == nil { 891 return 892 } 893 dst.or(src, runInline) 894 } 895 896 func (dst *Bitmap) or(src *Bitmap, runMode int) { 897 srcIdx, numKeys := 0, src.keys.numKeys() 898 899 buf := make([]uint16, maxContainerSize) 900 for ; srcIdx < numKeys; srcIdx++ { 901 srcCont := src.getContainer(src.keys.val(srcIdx)) 902 if getCardinality(srcCont) == 0 { 903 continue 904 } 905 906 key := src.keys.key(srcIdx) 907 908 dstIdx := dst.keys.search(key) 909 if dstIdx >= dst.keys.numKeys() || dst.keys.key(dstIdx) != key { 910 // srcCont doesn't exist in dst. So, copy it over. 911 offset := dst.newContainer(uint16(len(srcCont))) 912 copy(dst.getContainer(offset), srcCont) 913 dst.setKey(key, offset) 914 } else { 915 // Container exists in dst as well. Do an inline containerOr. 916 offset := dst.keys.val(dstIdx) 917 dstCont := dst.getContainer(offset) 918 if c := containerOr(dstCont, srcCont, buf, runMode|runInline); len(c) > 0 { 919 dst.copyAt(offset, c) 920 dst.setKey(key, offset) 921 } 922 } 923 } 924 } 925 926 func Or(a, b *Bitmap) *Bitmap { 927 ai, an := 0, a.keys.numKeys() 928 bi, bn := 0, b.keys.numKeys() 929 930 buf := make([]uint16, maxContainerSize) 931 res := NewBitmap() 932 for ai < an && bi < bn { 933 ak := a.keys.key(ai) 934 ac := a.getContainer(a.keys.val(ai)) 935 936 bk := b.keys.key(bi) 937 bc := b.getContainer(b.keys.val(bi)) 938 939 if ak == bk { 940 // Do the union. 941 outc := containerOr(ac, bc, buf, 0) 942 offset := res.newContainer(uint16(len(outc))) 943 copy(res.data[offset:], outc) 944 res.setKey(ak, offset) 945 ai++ 946 bi++ 947 } else if ak < bk { 948 off := res.newContainer(uint16(len(ac))) 949 copy(res.getContainer(off), ac) 950 res.setKey(ak, off) 951 ai++ 952 } else { 953 off := res.newContainer(uint16(len(bc))) 954 copy(res.getContainer(off), bc) 955 res.setKey(bk, off) 956 bi++ 957 } 958 } 959 for ai < an { 960 ak := a.keys.key(ai) 961 ac := a.getContainer(a.keys.val(ai)) 962 off := res.newContainer(uint16(len(ac))) 963 964 copy(res.getContainer(off), ac) 965 res.setKey(ak, off) 966 ai++ 967 } 968 for bi < bn { 969 bk := b.keys.key(bi) 970 bc := b.getContainer(b.keys.val(bi)) 971 off := res.newContainer(uint16(len(bc))) 972 973 copy(res.getContainer(off), bc) 974 res.setKey(bk, off) 975 bi++ 976 } 977 return res 978 } 979 980 func (ra *Bitmap) Rank(x uint64) int { 981 key := x & mask 982 offset, has := ra.keys.getValue(key) 983 if !has { 984 return -1 985 } 986 c := ra.getContainer(offset) 987 y := uint16(x) 988 989 // Find the rank within the container 990 var rank int 991 switch c[indexType] { 992 case typeArray: 993 rank = array(c).rank(y) 994 case typeBitmap: 995 rank = bitmap(c).rank(y) 996 } 997 if rank < 0 { 998 return -1 999 } 1000 1001 // Add up cardinalities of all the containers on the left of container containing x. 1002 n := ra.keys.numKeys() 1003 for i := 0; i < n; i++ { 1004 if ra.keys.key(i) == key { 1005 break 1006 } 1007 cont := ra.getContainer(ra.keys.val(i)) 1008 rank += getCardinality(cont) 1009 } 1010 return rank 1011 } 1012 1013 func (ra *Bitmap) Cleanup() { 1014 type interval struct { 1015 start uint64 1016 end uint64 1017 } 1018 1019 // Find the ranges that needs to be removed in the key space and the container space. Also, 1020 // start the iteration from idx = 1 because we never remove the 0 key. 1021 var keyIntervals, contIntervals []interval 1022 for idx := 1; idx < ra.keys.numKeys(); idx++ { 1023 off := ra.keys.val(idx) 1024 cont := ra.getContainer(off) 1025 if getCardinality(cont) == 0 { 1026 ko := uint64(keyOffset(idx)) 1027 contIntervals = append(contIntervals, interval{off, off + uint64(cont[indexSize])}) 1028 keyIntervals = append(keyIntervals, interval{4 * ko, 4 * (ko + 2)}) 1029 } 1030 } 1031 if len(contIntervals) == 0 { 1032 return 1033 } 1034 1035 merge := func(intervals []interval) []interval { 1036 assert(len(intervals) > 0) 1037 1038 // Merge the ranges in order to reduce scootLeft 1039 merged := []interval{intervals[0]} 1040 for _, ir := range intervals[1:] { 1041 last := merged[len(merged)-1] 1042 if ir.start == last.end { 1043 last.end = ir.end 1044 merged[len(merged)-1] = last 1045 continue 1046 } 1047 merged = append(merged, ir) 1048 } 1049 return merged 1050 } 1051 1052 // Key intervals are already sorted, but container intervals needs to be sorted because 1053 // they are always added in the end of the ra.data. 1054 sort.Slice(contIntervals, func(i, j int) bool { 1055 return contIntervals[i].start < contIntervals[j].start 1056 }) 1057 1058 contIntervals = merge(contIntervals) 1059 keyIntervals = merge(keyIntervals) 1060 1061 // Cleanup the containers. 1062 moved := uint64(0) 1063 for _, ir := range contIntervals { 1064 assert(ir.start >= moved) 1065 sz := ir.end - ir.start 1066 ra.scootLeft(ir.start-moved, sz) 1067 ra.keys.updateOffsets(ir.end-moved-1, sz, false) 1068 moved += sz 1069 } 1070 1071 // Cleanup the key space. 1072 moved = uint64(0) 1073 for _, ir := range keyIntervals { 1074 assert(ir.start >= moved) 1075 sz := ir.end - ir.start 1076 ra.scootLeft(ir.start-moved, sz) 1077 1078 // sz is in number of u16s, hence number of key-value removed is sz/8. 1079 ra.keys.setNumKeys(ra.keys.numKeys() - int(sz/8)) 1080 ra.keys.setNodeSize(ra.keys.size() - int(sz)) 1081 ra.keys = ra.keys[:len(ra.keys)-int(sz/4)] 1082 ra.keys.updateOffsets(ir.end-moved-1, sz, false) 1083 moved += sz 1084 } 1085 } 1086 1087 func FastAnd(bitmaps ...*Bitmap) *Bitmap { 1088 if len(bitmaps) == 0 { 1089 return NewBitmap() 1090 } 1091 b := bitmaps[0] 1092 for _, bm := range bitmaps[1:] { 1093 b.And(bm) 1094 } 1095 b.Cleanup() 1096 return b 1097 } 1098 1099 // FastParOr would group up bitmaps and call FastOr on them concurrently. It 1100 // would then merge the groups into final Bitmap. This approach is simpler and 1101 // faster than operating at a container level, because we can't operate on array 1102 // containers belonging to the same Bitmap concurrently because array containers 1103 // can expand, leaving no clear boundaries. 1104 // 1105 // If FastParOr is called with numGo=1, it just calls FastOr. 1106 // 1107 // Experiments with numGo=4 shows that FastParOr would be 2x the speed of 1108 // FastOr, but 4x the memory usage, even under 50% CPU usage. So, use wisely. 1109 func FastParOr(numGo int, bitmaps ...*Bitmap) *Bitmap { 1110 if numGo == 1 { 1111 return FastOr(bitmaps...) 1112 } 1113 width := max(len(bitmaps)/numGo, 3) 1114 1115 var wg sync.WaitGroup 1116 var res []*Bitmap 1117 for start := 0; start < len(bitmaps); start += width { 1118 end := min(start+width, len(bitmaps)) 1119 res = append(res, nil) // Make space for result. 1120 wg.Add(1) 1121 1122 go func(start, end int) { 1123 idx := start / width 1124 res[idx] = FastOr(bitmaps[start:end]...) 1125 wg.Done() 1126 }(start, end) 1127 } 1128 wg.Wait() 1129 return FastOr(res...) 1130 } 1131 1132 // FastOr would merge given Bitmaps into one Bitmap. This is faster than 1133 // doing an OR over the bitmaps iteratively. 1134 func FastOr(bitmaps ...*Bitmap) *Bitmap { 1135 if len(bitmaps) == 0 { 1136 return NewBitmap() 1137 } 1138 if len(bitmaps) == 1 { 1139 return bitmaps[0] 1140 } 1141 1142 // We first figure out the container distribution across the bitmaps. We do 1143 // that by looking at the key of the container, and the cardinality. We 1144 // assume the worst-case scenario where the union would result in a 1145 // cardinality (per container) of the sum of cardinalities of each of the 1146 // corresponding containers in other bitmaps. 1147 containers := make(map[uint64]int) 1148 for _, b := range bitmaps { 1149 for i := 0; i < b.keys.numKeys(); i++ { 1150 offset := b.keys.val(i) 1151 cont := b.getContainer(offset) 1152 card := getCardinality(cont) 1153 containers[b.keys.key(i)] += card 1154 } 1155 } 1156 1157 // We use the above information to pre-generate the destination Bitmap and 1158 // allocate container sizes based on the calculated cardinalities. 1159 // var sz int 1160 dst := NewBitmap() 1161 // First create the keys. We do this as a separate step, because keys are 1162 // the left most portion of the data array. Adding space there requires 1163 // moving a lot of pieces. 1164 for key, card := range containers { 1165 if card > 0 { 1166 dst.setKey(key, 0) 1167 } 1168 } 1169 1170 // Then create the bitmap containers. 1171 for key, card := range containers { 1172 if card >= 4096 { 1173 offset := dst.newContainer(maxContainerSize) 1174 c := dst.getContainer(offset) 1175 c[indexSize] = maxContainerSize 1176 c[indexType] = typeBitmap 1177 dst.setKey(key, offset) 1178 } 1179 } 1180 1181 // Create the array containers at the end. This allows them to expand 1182 // without having to move a lot of memory. 1183 for key, card := range containers { 1184 // Ensure this condition exactly maps up with above. 1185 if card < 4096 && card > 0 { 1186 if card < minContainerSize { 1187 card = minContainerSize 1188 } 1189 offset := dst.newContainer(uint16(card)) 1190 c := dst.getContainer(offset) 1191 c[indexSize] = uint16(card) 1192 c[indexType] = typeArray 1193 dst.setKey(key, offset) 1194 } 1195 } 1196 1197 // dst Bitmap is ready to be ORed with the given Bitmaps. 1198 for _, b := range bitmaps { 1199 dst.or(b, runLazy) 1200 } 1201 1202 for i := 0; i < dst.keys.numKeys(); i++ { 1203 offset := dst.keys.val(i) 1204 c := dst.getContainer(offset) 1205 if getCardinality(c) == invalidCardinality { 1206 calculateAndSetCardinality(c) 1207 } 1208 } 1209 1210 return dst 1211 } 1212 1213 // Split splits the bitmap based on maxSz and the externalSize function. It splits the bitmap 1214 // such that size of each split bitmap + external size corresponding to its elements approximately 1215 // equal to maxSz (it can be greater than maxSz sometimes). The splits are returned in sorted order. 1216 // externalSize is a function that should return the external size corresponding to elements in 1217 // range [start, end]. External size is used to calculate the split boundaries. 1218 func (bm *Bitmap) Split(externalSize func(start, end uint64) uint64, maxSz uint64) []*Bitmap { 1219 splitFurther := func(b *Bitmap) []*Bitmap { 1220 itr := b.NewIterator() 1221 newBm := NewBitmap() 1222 var sz uint64 1223 var bms []*Bitmap 1224 for id := itr.Next(); id != 0; id = itr.Next() { 1225 sz += externalSize(id, id) 1226 newBm.Set(id) 1227 if sz >= maxSz { 1228 bms = append(bms, newBm) 1229 newBm = NewBitmap() 1230 sz = 0 1231 } 1232 } 1233 1234 if !newBm.IsEmpty() { 1235 bms = append(bms, newBm) 1236 } 1237 return bms 1238 } 1239 1240 create := func(keyToOffset map[uint64]uint64, totalSz uint64) []*Bitmap { 1241 var keys []uint64 1242 for key := range keyToOffset { 1243 keys = append(keys, key) 1244 } 1245 sort.Slice(keys, func(i, j int) bool { 1246 return keys[i] < keys[j] 1247 }) 1248 1249 newBm := NewBitmap() 1250 1251 // First set all the keys. 1252 var containerSz uint64 1253 for _, key := range keys { 1254 newBm.setKey(key, 0) 1255 1256 // Calculate the size of the containers. 1257 cont := bm.getContainer(keyToOffset[key]) 1258 containerSz += uint64(len(cont)) 1259 } 1260 // Allocate enough space to hold all the containers. 1261 beforeSize := len(newBm.data) 1262 newBm.fastExpand(containerSz) 1263 newBm.data = newBm.data[:beforeSize] 1264 1265 // Now, we can populate the containers. For that, we first expand the 1266 // bitmap. Calculate the total size we need to allocate all these containers. 1267 for _, key := range keys { 1268 cont := bm.getContainer(keyToOffset[key]) 1269 off := newBm.newContainer(uint16(len(cont))) 1270 copy(newBm.data[off:], cont) 1271 1272 newBm.setKey(key, off) 1273 } 1274 1275 if newBm.GetCardinality() == 0 { 1276 return nil 1277 } 1278 1279 if totalSz > maxSz { 1280 return splitFurther(newBm) 1281 } 1282 1283 return []*Bitmap{newBm} 1284 } 1285 1286 var splits []*Bitmap 1287 1288 containerMap := make(map[uint64]uint64) 1289 var totalSz uint64 // size of containers plus the external size of the container 1290 1291 for i := 0; i < bm.keys.numKeys(); i++ { 1292 key := bm.keys.key(i) 1293 off := bm.keys.val(i) 1294 cont := bm.getContainer(off) 1295 1296 start, end := key, addUint64(key, 1<<16-1) 1297 sz := externalSize(start, end) + 2*uint64(cont[indexSize]) // Converting to bytes. 1298 1299 // We can probably append more containers in the same bucket. 1300 if totalSz+sz < maxSz || len(containerMap) == 0 { 1301 // Include this container in the container map. 1302 containerMap[key] = off 1303 totalSz += sz 1304 continue 1305 } 1306 1307 // We have reached the maxSz limit. Hence, create a split. 1308 splits = append(splits, create(containerMap, totalSz)...) 1309 1310 containerMap = make(map[uint64]uint64) 1311 containerMap[key] = off 1312 totalSz = sz 1313 } 1314 if len(containerMap) > 0 { 1315 splits = append(splits, create(containerMap, totalSz)...) 1316 } 1317 1318 return splits 1319 }