Merge pull request #1766 from richardartoul/ra/optimize-in-place-bulk

Implement fast roaring bitmap union (Part 1)
This commit is contained in:
Travis Turner 2018-12-07 15:52:23 -06:00 committed by GitHub
commit fef0b26c89
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 550 additions and 9 deletions

View file

@ -177,6 +177,15 @@ func (btc *bTreeContainers) Iterator(key uint64) (citer roaring.ContainerIterato
}, found
}
func (btc *bTreeContainers) Repair() {
e, _ := btc.tree.Seek(0)
_, c, err := e.Next()
for err != io.EOF {
c.Repair()
_, c, err = e.Next()
}
}
type btcIterator struct {
e *enumerator
key uint64

View file

@ -24,8 +24,8 @@ import (
"io/ioutil"
"net"
"net/http"
_ "net/http/pprof"
"net/url" // Imported for its side-effect of registering pprof endpoints with the server.
_ "net/http/pprof" // Imported for its side-effect of registering pprof endpoints with the server.
"net/url"
"reflect"
"runtime/debug"
"strconv"

View file

@ -156,6 +156,12 @@ func (sc *sliceContainers) Iterator(key uint64) (citer ContainerIterator, found
return &sliceIterator{e: sc, i: i}, found
}
func (sc *sliceContainers) Repair() {
for _, c := range sc.containers {
c.Repair()
}
}
type sliceIterator struct {
e *sliceContainers
i int

View file

@ -98,8 +98,12 @@ type Containers interface {
Count() uint64
//Reset will clear the containers collection to allow for recycling during snapshot
// Reset clears the containers collection to allow for recycling during snapshot
Reset()
// Repair will repair the cardinality of any containers whose cardinality were corrupted
// due to optimized operations.
Repair()
}
type ContainerIterator interface {
@ -395,10 +399,26 @@ func (b *Bitmap) Intersect(other *Bitmap) *Bitmap {
return output
}
// Union returns the bitwise union of b and other.
func (b *Bitmap) Union(other *Bitmap) *Bitmap {
// Union returns the bitwise union of b and others as a new bitmap.
func (b *Bitmap) Union(others ...*Bitmap) *Bitmap {
output := NewBitmap()
if len(others) == 1 {
b.unionIntoTargetSingle(output, others[0])
return output
}
output.UnionInPlace(b)
output.UnionInPlace(others...)
return output
}
// UnionInPlace returns the bitwise union of b and others, modifying
// b in place.
func (b *Bitmap) UnionInPlace(others ...*Bitmap) {
b.unionInPlace(others...)
}
func (b *Bitmap) unionIntoTargetSingle(target *Bitmap, other *Bitmap) {
iiter, _ := b.Containers.Iterator(0)
jiter, _ := other.Containers.Iterator(0)
i, j := iiter.Next(), jiter.Next()
@ -406,21 +426,259 @@ func (b *Bitmap) Union(other *Bitmap) *Bitmap {
kj, cj := jiter.Value()
for i || j {
if i && (!j || ki < kj) {
output.Containers.Put(ki, ci.Clone())
target.Containers.Put(ki, ci.Clone())
i = iiter.Next()
ki, ci = iiter.Value()
} else if j && (!i || ki > kj) {
output.Containers.Put(kj, cj.Clone())
target.Containers.Put(kj, cj.Clone())
j = jiter.Next()
kj, cj = jiter.Value()
} else { // ki == kj
output.Containers.Put(ki, union(ci, cj))
target.Containers.Put(ki, union(ci, cj))
i, j = iiter.Next(), jiter.Next()
ki, ci = iiter.Value()
kj, cj = jiter.Value()
}
}
return output
}
// unionIntoTarget stores the union of b and others into target. b and others will
// be left unchanged (unless one of them is also target), but target will be modified
// in place.
//
// This function performs an n-way union of n bitmaps. It performs this in an
// optimized manner looping through all the bitmaps and performing unions one
// container at a time. As a result, instead of generating many intermediary
// containers for each union operation for a given container key, only one
// new container needs to be allocated (or re-used) regardless of how many bitmaps
// participate in the union. This significantly reduces allocations. In addition,
// because we perform the unions one container at a time across all the bitmaps, we
// can calculate summary statistics that allow us to make more efficient decisions
// up front. For example, imagine trying to perform a union across the following three
// bitsets:
//
// 1. Bitmap A: Single array container at key 0 with 400 values in it.
// 2. Bitmap B: Single array container at key 0 with 500 values in it.
// 3. Bitmap C: Single array container at key 0 with 3500 values in it.
//
// Naive approach:
//
// 1. Perform union of bitmap A and B, container by container
// a. 400 + 500 < ArrayMaxSize so likely we will choose to allocate a new array
// container and then perform a unionArrayArray operation to merge the two
// arrays into the new array container.
// 2. Perform a union of the bitmap generated in the step above with bitmap C.
// 900 + 3500 > ArrayMaxSize so we will need to upgrade to a bitset container which
// we will have to allocate, and then we will have to perform two unions into the
// new bitmap container: one for the array container generated in the previous step,
// and one for the bitset container in bitmap C.
//
// Approach taken by this function:
//
// 1. Detect that bitmaps A, B, and C all have containers for key 0.
// 2. Estimate the resulting cardinality of the union of all their containers to be
// 400 + 500 + 3500 > ArrayMaxSize and decide upfront to use a bitset for the target
// container. Note that this is just an approximation of the final cardinality and can
// be off by a wide margin if there is a lot of overlap between containers, but that is
// fine, we'll still get the same result at the end, we'll just be more biased towards
// using bitmap containers will still being able to use array containers when all the
// cardinalities are small.
// 3. Union the containers from bitmaps A, B, and C into the new bitset container directly
// using fast bitwise operations.
//
// In the naive approach, we had to allocate two containers, whereas in the optimized approach
// we only had to allocate one container, and we also had to perform less union operations. This
// example is simplistic, but the impact in terms of CPU cycles and memory allocations achieved
// by using the optimized alogorithm when unioning many large bitmaps can be huge.
//
// An additional optimization that this function makes is that it recognizes that even when
// CPU support is present, performing the popcount() operation isn't free. Imagine a scenario
// where 10 bitset containers are being unioned together one after the next. If every
// bitset<->bitset union operation needs to keep the containers' cardinality up to date, then
// the algorithm will waste a lot of time performing intermediary popcount() operations that
// will immediately be invalidated by the next union operation. As a result, we allow the cardinality
// of containers to degrade when we perform the in-place union operations, and then when the algorithm
// completes we "repair" all the containers by performing the popcount() operation one time. This means
// that we only ever have to do O(1) popcount operations per container instead of O(n) where n is the
// number of containers with the same key that are being unioned together.
//
// The algorithm works by iterating through all of the containers in all of the bitmaps concurrently.
// At every "tick" of the outermost loop, we increment our pointer into the bitmaps list of containers
// by 1 (if we haven't reached the end of the containers for that bitmap.)
//
// We then loop through all of the "current" values(containers) for all of the bitmaps
// and for each container with a specific key that we encounter, we scan forward to see if any of the
// other bitmaps have a container for the same key. If so, we calculate some summary statistics and
// then use that information to make a decision about how to union all of the containers with the same
// key together, perform the union, mark the unioned containers as "handled" and then move on to the next
// batch of containers that share the same key.
//
// We repeat this process until every single bitmaps current container has been "handled". Then we start the
// outer loop over again and the process repeats until we've iterated through every container in every bitmap
// and unioned everything into a single target bitmap.
//
// The diagram below shows the iteration state of four different bitmaps as the algorithm progresses them.
// The diagrams should BE interpreted from left -> right, top -> bottom. The X's represent a container in
// the bitmap at a specific key, ^ symbol represents the bitmaps current container iteration position,
// and the - symbol represents a container that is at the current iteration position, but has been marked as "handled".
//
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 1 |___X____________X__________| | |___X____________X__________| | |___X____________X__________|
// ^ | _ |
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 2 |_______X________X______X___| | |_______X_______________X___| | |_______X_______________X___|
// ^ | ^ |
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 3 |_______X___________________| | |_______X___________________| | |_______X___________________|
// ^ | ^ |
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 4 |___X_______________________| | |___X_______________________| | |___X_______________________|
// ^ | _ |
// ------------------------------------------------------------------------------------------------------------------------
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 1 |___X____________X__________| | |___X____________X__________| | |___X____________X__________|
// _ | ^ | _
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 2 |_______X_______________X___| | |_______X_______________X___| | |_______X_______________X___|
// _ | ^ | ^
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 3 |_______X___________________| | |_______X___________________| | |_______X___________________|
// _ | |
// ---------------------------- | ---------------------------- | ----------------------------
// Bitmap 4 |___X_______________________| | |___X_______________________| | |___X_______________________|
// _
func (b *Bitmap) unionInPlace(others ...*Bitmap) {
var (
requiredSliceSize = len(others)
// To avoid having to allocate a slice everytime, if the number of bitmaps
// being unioned is small enough we can just use this stack-allocated array.
staticHandledIters = [20]handledIter{}
bitmapIters handledIters
target = b
)
if requiredSliceSize <= 20 {
bitmapIters = staticHandledIters[:0]
} else {
bitmapIters = make(handledIters, 0, requiredSliceSize)
}
for _, other := range others {
otherIter, _ := other.Containers.Iterator(0)
if otherIter.Next() {
bitmapIters = append(bitmapIters, handledIter{
iter: otherIter,
hasNext: true,
handled: false,
})
}
}
// Loop until we've exhausted every iter.
hasNext := true
for hasNext {
// Loop until every iters current value has been handled.
for i, iIter := range bitmapIters {
if !iIter.hasNext || iIter.handled {
// Either we've exhausted this iter (it has no more containers), or
// we've already handled the current container by unioning it with
// one of the containers we encountered earlier.
continue
}
var (
iKey, iContainer = iIter.iter.Value()
// Summary statistics about all the containers in the other bitmaps
// that share the same key so we can make smarter union strategy
// decisions later. Note that we slice to [i:] not [i+1:] because we
// want to include the current containers information in the stats.
summaryStats = bitmapIters[i:].calculateSummaryStats(iKey)
)
if summaryStats.isOnlyContainerWithKey {
// TODO(rartoul): We can avoid these clones if we can determine
// if the container is coming from an immutable bitmap and we
// know that we can mark the target bitmap as immutable as well.
target.Containers.Put(iKey, iContainer.Clone())
bitmapIters[i].handled = true
continue
}
// There was more than one container across the bitmaps with key iKey
// so we need to calculate a union.
if summaryStats.hasMaxRange {
// One (or more) of the containers represented the maximum possible
// range that a container can store, so instead of calculating a
// union we can generate an RLE container that represents the entire
// range.
container := &Container{
runs: []interval16{{start: 0, last: maxContainerVal}},
containerType: containerRun,
n: maxContainerVal + 1,
}
target.Containers.Put(iKey, container)
bitmapIters[i:].markItersWithCurrentKeyAsHandled(iKey)
continue
}
// Use a bitmap container for the target bitmap, and union everything
// into it.
//
// TODO(rartoul): Add another conditional case for n < ArrayMaxSize
// (or some fraction of that value) to avoid allocating expensive
// bitmaps when unioning many low-density array containers, but this
// will require writing a union in place algorithm for an array container
// that accepts multiple different containers to union into it for
// efficiency.
container := target.Containers.Get(iKey)
// If target already has a bitmap container for iKey then we can reuse that,
// otherwise we have to allocate a new one.
if container == nil || container.containerType != containerBitmap {
buf := make([]uint64, bitmapN)
ob := buf[:bitmapN]
container = &Container{
bitmap: ob,
n: 0,
containerType: containerBitmap,
}
}
// Once we've acquired a bitmap container (either by reusing the existing one
// or allocating a new one) then the last step is to iterate through all the
// other containers to see which ones have the same key, and union all of them
// into the target bitmap container. Only need to loop starting from i because
// anything previous to that has already been handled.
for j, jIter := range bitmapIters[i:] {
jKey, jContainer := jIter.iter.Value()
if iKey == jKey {
if jContainer.isArray() {
unionBitmapArrayInPlace(container, jContainer)
} else if jContainer.isRun() {
unionBitmapRunInPlace(container, jContainer)
} else {
unionBitmapBitmapInPlace(container, jContainer)
}
bitmapIters[j].handled = true
}
}
// Now that we've calculated a container that is a union of all the containers
// with the same key across all the bitmaps, we store it in the list of containers
// for the target.
target.Containers.Put(iKey, container)
}
hasNext = bitmapIters.next()
}
// Performing the popcount() operation with every union is wasteful because
// its likely the value will be invalidated by the next union operation. As
// a result, when we're performing all of our in-place unions we allow the value of
// n (container cardinality) to fall out of sync, and then at the very end we perform
// a "Repair" to recalculate all the container values. That way we never popcount()
// an entire bitmap container more than once per bulk union operation.
target.Containers.Repair()
}
// Difference returns the difference of b and other.
@ -1824,6 +2082,28 @@ func (c *Container) check() error {
return a
}
// Repair repairs the cardinality of c if it has been corrupted by
// optimized operations.
func (c *Container) Repair() {
if c.isBitmap() {
c.bitmapRepair()
}
}
func (c *Container) bitmapRepair() {
n := int32(0)
// Manually unroll loop to make it a little faster.
// TODO(rartoul): Can probably make this a few x faster using
// SIMD instructions.
for i := 0; i < bitmapN; i += 4 {
n += int32(popcount(c.bitmap[i]))
n += int32(popcount(c.bitmap[i+1]))
n += int32(popcount(c.bitmap[i+2]))
n += int32(popcount(c.bitmap[i+3]))
}
c.n = n
}
// containerInfo represents a point-in-time snapshot of container stats.
type containerInfo struct {
Key uint64 // container key
@ -2384,6 +2664,15 @@ func unionBitmapRun(a, b *Container) *Container {
return output
}
// unions the run b into the bitmap a, mutating a in place. The n value of
// a will need to be repaired after the fact.
func unionBitmapRunInPlace(a, b *Container) {
statsHit("union/BitmapRun")
for j := 0; j < len(b.runs); j++ {
a.bitmapSetRangeIgnoreN(uint64(b.runs[j].start), uint64(b.runs[j].last)+1)
}
}
const maxBitmap = 0xFFFFFFFFFFFFFFFF
// sets all bits in [i, j) (c must be a bitmap container)
@ -2409,6 +2698,25 @@ func (c *Container) bitmapSetRange(i, j uint64) {
}
}
// sets all bits in [i, j) (c must be a bitmap container) without updating
// the value of n, meaning it will need to be repaired after the fact.
func (c *Container) bitmapSetRangeIgnoreN(i, j uint64) {
x := i >> 6
y := (j - 1) >> 6
var X uint64 = maxBitmap << (i % 64)
var Y uint64 = maxBitmap >> (63 - ((j - 1) % 64))
if x == y {
c.bitmap[x] |= (X & Y)
} else {
c.bitmap[x] |= X
for i := x + 1; i < y; i++ {
c.bitmap[i] = maxBitmap
}
c.bitmap[y] |= Y
}
}
// xor's all bits in [i, j) with all true (c must be a bitmap container).
func (c *Container) bitmapXorRange(i, j uint64) {
x := i >> 6
@ -2503,6 +2811,14 @@ func unionArrayBitmap(a, b *Container) *Container {
return output
}
// unions array b into bitmap a, mutating a in place. The n value
// of a will need to be repaired after the fact.
func unionBitmapArrayInPlace(a, b *Container) {
for _, v := range b.array {
a.bitmap[v>>6] |= (uint64(1) << (v % 64))
}
}
func unionBitmapBitmap(a, b *Container) *Container {
// local variables added to prevent BCE checks in loop
// see https://go101.org/article/bounds-check-elimination.html
@ -2529,6 +2845,28 @@ func unionBitmapBitmap(a, b *Container) *Container {
return output
}
// unions bitmap b into bitmap a, mutating a in place. The n value of
// a will need to be repaired after the fact.
func unionBitmapBitmapInPlace(a, b *Container) {
// local variables added to prevent BCE checks in loop
// see https://go101.org/article/bounds-check-elimination.html
var (
ab = a.bitmap[:bitmapN]
bb = b.bitmap[:bitmapN]
)
// Manually unroll loop to make it a little faster.
// TODO(rartoul): Can probably make this a few x faster using
// SIMD instructions.
for i := 0; i < bitmapN; i += 4 {
ab[i] |= bb[i]
ab[i+1] |= bb[i+1]
ab[i+2] |= bb[i+2]
ab[i+3] |= bb[i+3]
}
}
func difference(a, b *Container) *Container {
if a.isArray() {
if b.isArray() {
@ -3637,3 +3975,82 @@ func readWithRuns(b *Bitmap, data []byte, pos int, keyN uint32) {
}
}
}
// handledIter and handledIters are wrappers around Bitmap Container iterators
// and assist with the unionIntoTarget algorithm by abstracting away some tedious
// operations.
type handledIter struct {
iter ContainerIterator
hasNext bool
handled bool
}
type handledIters []handledIter
func (w handledIters) next() bool {
hasNext := false
for i, wrapped := range w {
next := wrapped.iter.Next()
w[i].hasNext = next
w[i].handled = false
if next {
hasNext = true
}
}
return hasNext
}
func (w handledIters) markItersWithCurrentKeyAsHandled(key uint64) {
for i, wrapped := range w {
currKey, _ := wrapped.iter.Value()
if currKey == key {
w[i].handled = true
}
}
}
func (w handledIters) calculateSummaryStats(key uint64) containerUnionSummaryStats {
summary := containerUnionSummaryStats{}
for _, iter := range w {
// Calculate key-level statistics here
currKey, currContainer := iter.iter.Value()
if key == currKey {
summary.isOnlyContainerWithKey = false
summary.n += currContainer.n
if currContainer.n == maxContainerVal+1 {
summary.hasMaxRange = true
summary.n = maxContainerVal + 1
return summary
}
}
}
return summary
}
// Summary statistics about all the containers in the other bitmaps
// that share the same key so we can make smarter union strategy
// decisions.
type containerUnionSummaryStats struct {
// Estimated cardinality of the union of all containers with the same
// key across all bitmaps. This calculation is very rough as we just sum
// the cardinality of the container across the different bitmaps which could
// result in very inflated values, but it allows us to avoid allocating
// expensive bitmaps when unioning many low density containers.
n int32
// Whether any other is the only container across all the bitmaps
// with the specified key. If true, we can skip all the unioning logic
// and just clone the container into target.
isOnlyContainerWithKey bool
// Whether any of the containers with the specified keys are storing every possible
// value that they can. If so, we can short-circuit all the unioning logic and use
// a RLE container with a single value in it. This is an optimization to
// avoid using an expensive bitmap container for bitmaps that have some
// extremely dense containers.
hasMaxRange bool
}

View file

@ -3267,6 +3267,34 @@ func TestUnmarshalOfficialRoaring(t *testing.T) {
}
func BenchmarkUnionBitmapBitmapInPlace(b *testing.B) {
b1 := newTestBitmapContainer()
b2 := newTestBitmapContainer()
for n := 0; n < b.N; n++ {
unionBitmapBitmapInPlace(b1, b2)
}
}
func BenchmarkBitmapRepair(b *testing.B) {
b1 := newTestBitmapContainer()
for n := 0; n < b.N; n++ {
b1.bitmapRepair()
}
}
func newTestBitmapContainer() *Container {
var (
buf = make([]uint64, bitmapN)
ob = buf[:bitmapN]
container = &Container{
bitmap: ob,
n: 0,
containerType: containerBitmap,
}
)
return container
}
/*
// This function exercises an arcane edge case in dead code.
// It doesn't need to be run right now.

View file

@ -396,7 +396,44 @@ func TestBitmap_Union1(t *testing.T) {
if n := result.Count(); n != 75007 {
t.Fatalf("unexpected n: %d", n)
}
}
func TestBitmap_UnionInPlace1(t *testing.T) {
var (
bm0 = roaring.NewFileBitmap(0, 2683177)
bm1 = roaring.NewFileBitmap()
result = roaring.NewBitmap()
)
for i := uint64(628); i < 2683301; i++ {
bm1.Add(i)
}
bm1.Add(4000000)
result.UnionInPlace(bm0, bm1)
if n := result.Count(); n != 2682675 {
t.Fatalf("unexpected n: %d", n)
}
bm := testBM()
result = roaring.NewBitmap()
result.UnionInPlace(bm, bm0)
if n := result.Count(); n != 75009 {
t.Fatalf("unexpected n: %d", n)
}
result = roaring.NewBitmap()
result.UnionInPlace(bm, bm)
if n := result.Count(); n != 75007 {
t.Fatalf("unexpected n: %d", n)
}
// Make sure the bitmaps weren't mutated.
if n := bm0.Count(); n != 2 {
t.Fatalf("unexpected n: %d", n)
}
if n := bm1.Count(); n != 2682674 {
t.Fatalf("unexpected n: %d", n)
}
}
func TestBitmap_Intersection_Empty(t *testing.T) {
@ -544,11 +581,35 @@ func TestBitmap_Union(t *testing.T) {
bm0 := roaring.NewFileBitmap(0, 1000001, 1000002, 1000003)
bm1 := roaring.NewFileBitmap(0, 50000, 1000001, 1000002)
result := bm0.Union(bm1)
if n := result.Count(); n != 5 {
t.Fatalf("unexpected n: %d", n)
}
}
func TestBitmap_UnionInPlace(t *testing.T) {
var (
bm0 = roaring.NewFileBitmap(0, 1000001, 1000002, 1000003)
bm1 = roaring.NewFileBitmap(0, 50000, 1000001, 1000002)
result = roaring.NewBitmap()
)
result.UnionInPlace(bm0, bm1)
// Make sure the union worked.
if n := result.Count(); n != 5 {
t.Fatalf("unexpected n: %d", n)
}
// Make sure the other bitmaps weren't mutated.
if n := bm0.Count(); n != 4 {
t.Fatalf("unexpected n: %d", n)
}
if n := bm1.Count(); n != 4 {
t.Fatalf("unexpected n: %d", n)
}
}
func TestBitmap_Xor(t *testing.T) {
bm0 := testBM()
bm1 := roaring.NewFileBitmap(0, 1, 2, 3)
@ -1350,3 +1411,23 @@ func BenchmarkSliceDescending(b *testing.B) {
}
}
}
func BenchmarkUnion(b *testing.B) {
data := getBenchData(b)
for n := 0; n < b.N; n++ {
data.a1.
Union(data.a2).
Union(data.b).
Union(data.r1).
Union(data.r2)
}
}
func BenchmarkUnionBulk(b *testing.B) {
data := getBenchData(b)
bm := roaring.NewBitmap()
for n := 0; n < b.N; n++ {
bm.
UnionInPlace(data.a1, data.a2, data.b, data.r1, data.r2)
}
}