From 57bb51c45852b0993f95b939121b5a32ee9eaf80 Mon Sep 17 00:00:00 2001 From: Kevin Conaway Date: Sun, 29 Mar 2020 08:41:39 -0400 Subject: [PATCH 1/3] Add Iterate() method for iterating the bitmap without an iterator One of our applications currently makes heavy use of Iterators. We noticed that they are a minor but noticeable source of allocations (along with the associated cpu & gc cost) in our application due to each Iterator being heap allocated along with the internal iterators each container creates. This commit provides an alternate Iterate() method with a callback that iterates the bitmap in place, invoking the callback for each value. This avoids allocations entirely at the cost of some code duplication --- arraycontainer.go | 6 ++++ benchmark_test.go | 89 ++++++++++++++++++++++++++++++++++++++-------- bitmapcontainer.go | 8 +++++ roaring.go | 28 +++++++++++++++ roaring_test.go | 48 +++++++++++++++++++++++++ roaringarray.go | 4 ++- runcontainer.go | 24 +++++++++++++ 7 files changed, 191 insertions(+), 16 deletions(-) diff --git a/arraycontainer.go b/arraycontainer.go index 621616f5..b227f9f3 100644 --- a/arraycontainer.go +++ b/arraycontainer.go @@ -24,6 +24,12 @@ func (ac *arrayContainer) fillLeastSignificant16bits(x []uint32, i int, mask uin } } +func (ac *arrayContainer) iterate(cb func(x uint16)) { + for i := 0; i < len(ac.content); i++ { + cb(ac.content[i]) + } +} + func (ac *arrayContainer) getShortIterator() shortPeekable { return &shortIterator{ac.content, 0} } diff --git a/benchmark_test.go b/benchmark_test.go index b90a5fe5..0df917b8 100644 --- a/benchmark_test.go +++ b/benchmark_test.go @@ -363,23 +363,82 @@ func BenchmarkCountBitset(b *testing.B) { // go test -bench BenchmarkIterate -run - func BenchmarkIterateRoaring(b *testing.B) { - b.StopTimer() - r := rand.New(rand.NewSource(0)) - s := NewBitmap() - sz := 150000 - initsize := 65000 - for i := 0; i < initsize; i++ { - s.Add(uint32(r.Int31n(int32(sz)))) - } - b.StartTimer() - for j := 0; j < b.N; j++ { - c9 = uint(0) - i := s.Iterator() - for i.HasNext() { - i.Next() - c9++ + newBitmap := func() *Bitmap { + r := rand.New(rand.NewSource(0)) + s := NewBitmap() + sz := 150000 + initsize := 65000 + for i := 0; i < initsize; i++ { + s.Add(uint32(r.Int31n(int32(sz)))) } + return s } + + b.Run("iterator-compressed", func(b *testing.B) { + b.ReportAllocs() + + s := newBitmap() + s.RunOptimize() + + b.ResetTimer() + + for j := 0; j < b.N; j++ { + c9 = uint(0) + i := s.Iterator() + for i.HasNext() { + i.Next() + c9++ + } + } + }) + + b.Run("iterator", func(b *testing.B) { + b.ReportAllocs() + + s := newBitmap() + + b.ResetTimer() + + for j := 0; j < b.N; j++ { + c9 = uint(0) + i := s.Iterator() + for i.HasNext() { + i.Next() + c9++ + } + } + }) + + b.Run("iterate-compressed", func(b *testing.B) { + b.ReportAllocs() + + s := newBitmap() + s.RunOptimize() + + b.ResetTimer() + + for j := 0; j < b.N; j++ { + c9 = uint(0) + s.Iterate(func(x uint32) { + c9++ + }) + } + }) + + b.Run("iterate", func(b *testing.B) { + b.ReportAllocs() + + s := newBitmap() + + b.ResetTimer() + + for j := 0; j < b.N; j++ { + c9 = uint(0) + s.Iterate(func(x uint32) { + c9++ + }) + } + }) } // go test -bench BenchmarkSparseIterate -run - diff --git a/bitmapcontainer.go b/bitmapcontainer.go index e749721b..44f75305 100644 --- a/bitmapcontainer.go +++ b/bitmapcontainer.go @@ -96,6 +96,14 @@ func (bc *bitmapContainer) maximum() uint16 { return uint16(0) } +func (bc *bitmapContainer) iterate(cb func(x uint16)) { + for i := bc.NextSetBit(0); i >= 0; { + j := i + i = bc.NextSetBit(i + 1) + cb(uint16(j)) + } +} + type bitmapContainerShortIterator struct { ptr *bitmapContainer i int diff --git a/roaring.go b/roaring.go index 51724d67..58ade9dc 100644 --- a/roaring.go +++ b/roaring.go @@ -416,6 +416,34 @@ func (rb *Bitmap) String() string { return buffer.String() } +// Iterate iterates over the bitmap, calling the given callback with each value in the bitmap. +// The iteration results are undefined if the bitmap is modified (e.g., with Add or Remove). +// There is no guarantee as to what order the values will be iterated +func (rb *Bitmap) Iterate(cb func(x uint32)) { + var hs uint32 + for i := 0; i < rb.highlowcontainer.size(); i++ { + hs = uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16 + + c := rb.highlowcontainer.getContainerAtIndex(i) + + // This is hacky but it avoids allocations from invoking an interface method with a closure + switch t := c.(type) { + case *arrayContainer: + t.iterate(func(x uint16) { + cb(uint32(x) | hs) + }) + case *runContainer16: + t.iterate(func(x uint16) { + cb(uint32(x) | hs) + }) + case *bitmapContainer: + t.iterate(func(x uint16) { + cb(uint32(x) | hs) + }) + } + } +} + // Iterator creates a new IntPeekable to iterate over the integers contained in the bitmap, in sorted order; // the iterator becomes invalid if the bitmap is modified (e.g., with Add or Remove). func (rb *Bitmap) Iterator() IntPeekable { diff --git a/roaring_test.go b/roaring_test.go index 5a031a0d..68e4ff36 100644 --- a/roaring_test.go +++ b/roaring_test.go @@ -2308,3 +2308,51 @@ func TestBitmapFlipMaxRangeEnd(t *testing.T) { assert.EqualValues(t, MaxRange, bm.GetCardinality()) } + +func TestIterate(t *testing.T) { + rb := NewBitmap() + + for i := 0; i < 300; i++ { + rb.Add(uint32(i)) + } + + var values []uint32 + rb.Iterate(func(x uint32) { + values = append(values, x) + }) + + assert.Equal(t, rb.ToArray(), values) +} + +func TestIterateCompressed(t *testing.T) { + rb := NewBitmap() + + for i := 0; i < 300; i++ { + rb.Add(uint32(i)) + } + + rb.RunOptimize() + + var values []uint32 + rb.Iterate(func(x uint32) { + values = append(values, x) + }) + + assert.Equal(t, rb.ToArray(), values) +} + +func TestIterateLargeValues(t *testing.T) { + rb := NewBitmap() + + // This range of values ensures that all different types of containers will be used + for i := 150000; i < 450000; i++ { + rb.Add(uint32(i)) + } + + var values []uint32 + rb.Iterate(func(x uint32) { + values = append(values, x) + }) + + assert.Equal(t, rb.ToArray(), values) +} diff --git a/roaringarray.go b/roaringarray.go index d9d5edda..312b7627 100644 --- a/roaringarray.go +++ b/roaringarray.go @@ -4,9 +4,10 @@ import ( "bytes" "encoding/binary" "fmt" + "io" + snappy "github.com/glycerine/go-unsnap-stream" "github.com/tinylib/msgp/msgp" - "io" ) //go:generate msgp -unexported @@ -38,6 +39,7 @@ type container interface { inot(firstOfRange, endx int) container // i stands for inplace, range is [firstOfRange,endx) xor(r container) container getShortIterator() shortPeekable + iterate(cb func(x uint16)) getReverseIterator() shortIterable getManyIterator() manyIterable contains(i uint16) bool diff --git a/runcontainer.go b/runcontainer.go index cbffdaf2..27770643 100644 --- a/runcontainer.go +++ b/runcontainer.go @@ -1162,6 +1162,30 @@ func (rc *runContainer16) newRunIterator16() *runIterator16 { return &runIterator16{rc: rc, curIndex: 0, curPosInIndex: 0} } +func (rc *runContainer16) iterate(cb func(x uint16)) { + curIndex := int64(0) + curPosInIndex := uint16(0) + + hasNext := func() bool { + return int64(len(rc.iv)) > curIndex+1 || + (int64(len(rc.iv)) == curIndex+1 && rc.iv[curIndex].length >= curPosInIndex) + } + + for hasNext() { + next := rc.iv[curIndex].start + curPosInIndex + + if curPosInIndex == rc.iv[curIndex].length { + curPosInIndex = 0 + curIndex++ + } else { + curPosInIndex++ + } + + cb(next) + } + +} + // hasNext returns false if calling next will panic. It // returns true when there is at least one more value // available in the iteration sequence. From 06686dbc23a53ef9cdbc18f860a2ebf106c98f02 Mon Sep 17 00:00:00 2001 From: Kevin Conaway Date: Sun, 29 Mar 2020 10:28:30 -0400 Subject: [PATCH 2/3] Add support for halting the iteration --- arraycontainer.go | 8 ++++++-- benchmark_test.go | 6 ++++-- bitmapcontainer.go | 7 +++++-- roaring.go | 21 +++++++++++++-------- roaring_test.go | 34 +++++++++++++++++++++++++++++++--- roaringarray.go | 2 +- runcontainer.go | 7 +++++-- 7 files changed, 65 insertions(+), 20 deletions(-) diff --git a/arraycontainer.go b/arraycontainer.go index b227f9f3..8001b994 100644 --- a/arraycontainer.go +++ b/arraycontainer.go @@ -24,10 +24,14 @@ func (ac *arrayContainer) fillLeastSignificant16bits(x []uint32, i int, mask uin } } -func (ac *arrayContainer) iterate(cb func(x uint16)) { +func (ac *arrayContainer) iterate(cb func(x uint16) bool) bool { for i := 0; i < len(ac.content); i++ { - cb(ac.content[i]) + if !cb(ac.content[i]) { + return false + } } + + return true } func (ac *arrayContainer) getShortIterator() shortPeekable { diff --git a/benchmark_test.go b/benchmark_test.go index 0df917b8..6693315a 100644 --- a/benchmark_test.go +++ b/benchmark_test.go @@ -419,8 +419,9 @@ func BenchmarkIterateRoaring(b *testing.B) { for j := 0; j < b.N; j++ { c9 = uint(0) - s.Iterate(func(x uint32) { + s.Iterate(func(x uint32) bool { c9++ + return true }) } }) @@ -434,8 +435,9 @@ func BenchmarkIterateRoaring(b *testing.B) { for j := 0; j < b.N; j++ { c9 = uint(0) - s.Iterate(func(x uint32) { + s.Iterate(func(x uint32) bool { c9++ + return true }) } }) diff --git a/bitmapcontainer.go b/bitmapcontainer.go index 44f75305..ea419a14 100644 --- a/bitmapcontainer.go +++ b/bitmapcontainer.go @@ -96,12 +96,15 @@ func (bc *bitmapContainer) maximum() uint16 { return uint16(0) } -func (bc *bitmapContainer) iterate(cb func(x uint16)) { +func (bc *bitmapContainer) iterate(cb func(x uint16) bool) bool { for i := bc.NextSetBit(0); i >= 0; { j := i i = bc.NextSetBit(i + 1) - cb(uint16(j)) + if !cb(uint16(j)) { + return false + } } + return true } type bitmapContainerShortIterator struct { diff --git a/roaring.go b/roaring.go index 58ade9dc..a1685447 100644 --- a/roaring.go +++ b/roaring.go @@ -416,31 +416,36 @@ func (rb *Bitmap) String() string { return buffer.String() } -// Iterate iterates over the bitmap, calling the given callback with each value in the bitmap. +// Iterate iterates over the bitmap, calling the given callback with each value in the bitmap. If the callback returns +// false, the iteration is halted. // The iteration results are undefined if the bitmap is modified (e.g., with Add or Remove). // There is no guarantee as to what order the values will be iterated -func (rb *Bitmap) Iterate(cb func(x uint32)) { +func (rb *Bitmap) Iterate(cb func(x uint32) bool) { var hs uint32 for i := 0; i < rb.highlowcontainer.size(); i++ { hs = uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16 c := rb.highlowcontainer.getContainerAtIndex(i) + var shouldContinue bool // This is hacky but it avoids allocations from invoking an interface method with a closure switch t := c.(type) { case *arrayContainer: - t.iterate(func(x uint16) { - cb(uint32(x) | hs) + shouldContinue = t.iterate(func(x uint16) bool { + return cb(uint32(x) | hs) }) case *runContainer16: - t.iterate(func(x uint16) { - cb(uint32(x) | hs) + shouldContinue = t.iterate(func(x uint16) bool { + return cb(uint32(x) | hs) }) case *bitmapContainer: - t.iterate(func(x uint16) { - cb(uint32(x) | hs) + shouldContinue = t.iterate(func(x uint16) bool { + return cb(uint32(x) | hs) }) } + if !shouldContinue { + return + } } } diff --git a/roaring_test.go b/roaring_test.go index 68e4ff36..bb0ece29 100644 --- a/roaring_test.go +++ b/roaring_test.go @@ -2317,8 +2317,9 @@ func TestIterate(t *testing.T) { } var values []uint32 - rb.Iterate(func(x uint32) { + rb.Iterate(func(x uint32) bool { values = append(values, x) + return true }) assert.Equal(t, rb.ToArray(), values) @@ -2334,8 +2335,9 @@ func TestIterateCompressed(t *testing.T) { rb.RunOptimize() var values []uint32 - rb.Iterate(func(x uint32) { + rb.Iterate(func(x uint32) bool { values = append(values, x) + return true }) assert.Equal(t, rb.ToArray(), values) @@ -2350,9 +2352,35 @@ func TestIterateLargeValues(t *testing.T) { } var values []uint32 - rb.Iterate(func(x uint32) { + rb.Iterate(func(x uint32) bool { values = append(values, x) + return true }) assert.Equal(t, rb.ToArray(), values) } + +func TestIterateHalt(t *testing.T) { + rb := NewBitmap() + + // This range of values ensures that all different types of containers will be used + for i := 150000; i < 450000; i++ { + rb.Add(uint32(i)) + } + + var values []uint32 + count := uint64(0) + stopAt := rb.GetCardinality() - 1 + rb.Iterate(func(x uint32) bool { + values = append(values, x) + count++ + if count == stopAt { + return false + } + return true + }) + + expected := rb.ToArray() + expected = expected[0 : len(expected)-1] + assert.Equal(t, expected, values) +} diff --git a/roaringarray.go b/roaringarray.go index 312b7627..6e35ac3e 100644 --- a/roaringarray.go +++ b/roaringarray.go @@ -39,7 +39,7 @@ type container interface { inot(firstOfRange, endx int) container // i stands for inplace, range is [firstOfRange,endx) xor(r container) container getShortIterator() shortPeekable - iterate(cb func(x uint16)) + iterate(cb func(x uint16) bool) bool getReverseIterator() shortIterable getManyIterator() manyIterable contains(i uint16) bool diff --git a/runcontainer.go b/runcontainer.go index 27770643..4da49c16 100644 --- a/runcontainer.go +++ b/runcontainer.go @@ -1162,7 +1162,7 @@ func (rc *runContainer16) newRunIterator16() *runIterator16 { return &runIterator16{rc: rc, curIndex: 0, curPosInIndex: 0} } -func (rc *runContainer16) iterate(cb func(x uint16)) { +func (rc *runContainer16) iterate(cb func(x uint16) bool) bool { curIndex := int64(0) curPosInIndex := uint16(0) @@ -1181,9 +1181,12 @@ func (rc *runContainer16) iterate(cb func(x uint16)) { curPosInIndex++ } - cb(next) + if !cb(next) { + return false + } } + return true } // hasNext returns false if calling next will panic. It From 7c87cea09c7676dcb85db16f99362f5db8456ff9 Mon Sep 17 00:00:00 2001 From: Kevin Conaway Date: Mon, 30 Mar 2020 11:54:39 -0400 Subject: [PATCH 3/3] Use stack allocated iterators in each container type to avoid duplicating iterator logic --- arraycontainer.go | 6 ++++-- bitmapcontainer.go | 9 +++++---- roaring.go | 7 +++---- runcontainer.go | 21 +++------------------ 4 files changed, 15 insertions(+), 28 deletions(-) diff --git a/arraycontainer.go b/arraycontainer.go index 8001b994..eb124f3b 100644 --- a/arraycontainer.go +++ b/arraycontainer.go @@ -25,8 +25,10 @@ func (ac *arrayContainer) fillLeastSignificant16bits(x []uint32, i int, mask uin } func (ac *arrayContainer) iterate(cb func(x uint16) bool) bool { - for i := 0; i < len(ac.content); i++ { - if !cb(ac.content[i]) { + iterator := shortIterator{ac.content, 0} + + for iterator.hasNext() { + if !cb(iterator.next()) { return false } } diff --git a/bitmapcontainer.go b/bitmapcontainer.go index ea419a14..cd259fd2 100644 --- a/bitmapcontainer.go +++ b/bitmapcontainer.go @@ -97,13 +97,14 @@ func (bc *bitmapContainer) maximum() uint16 { } func (bc *bitmapContainer) iterate(cb func(x uint16) bool) bool { - for i := bc.NextSetBit(0); i >= 0; { - j := i - i = bc.NextSetBit(i + 1) - if !cb(uint16(j)) { + iterator := bitmapContainerShortIterator{bc, bc.NextSetBit(0)} + + for iterator.hasNext() { + if !cb(iterator.next()) { return false } } + return true } diff --git a/roaring.go b/roaring.go index a1685447..ed75d58b 100644 --- a/roaring.go +++ b/roaring.go @@ -421,10 +421,8 @@ func (rb *Bitmap) String() string { // The iteration results are undefined if the bitmap is modified (e.g., with Add or Remove). // There is no guarantee as to what order the values will be iterated func (rb *Bitmap) Iterate(cb func(x uint32) bool) { - var hs uint32 for i := 0; i < rb.highlowcontainer.size(); i++ { - hs = uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16 - + hs := uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16 c := rb.highlowcontainer.getContainerAtIndex(i) var shouldContinue bool @@ -443,8 +441,9 @@ func (rb *Bitmap) Iterate(cb func(x uint32) bool) { return cb(uint32(x) | hs) }) } + if !shouldContinue { - return + break } } } diff --git a/runcontainer.go b/runcontainer.go index 4da49c16..5a0f985f 100644 --- a/runcontainer.go +++ b/runcontainer.go @@ -1163,25 +1163,10 @@ func (rc *runContainer16) newRunIterator16() *runIterator16 { } func (rc *runContainer16) iterate(cb func(x uint16) bool) bool { - curIndex := int64(0) - curPosInIndex := uint16(0) + iterator := runIterator16{rc, 0, 0} - hasNext := func() bool { - return int64(len(rc.iv)) > curIndex+1 || - (int64(len(rc.iv)) == curIndex+1 && rc.iv[curIndex].length >= curPosInIndex) - } - - for hasNext() { - next := rc.iv[curIndex].start + curPosInIndex - - if curPosInIndex == rc.iv[curIndex].length { - curPosInIndex = 0 - curIndex++ - } else { - curPosInIndex++ - } - - if !cb(next) { + for iterator.hasNext() { + if !cb(iterator.next()) { return false } }