diff --git a/fragment.go b/fragment.go index 5458ee34a..db0d3fb13 100644 --- a/fragment.go +++ b/fragment.go @@ -274,7 +274,7 @@ func (f *Fragment) openCache() error { // Read in all rows by ID. // This will cause them to be added to the cache. for _, id := range pb.IDs { - n := f.row(id, false, false).Count() + n := f.storage.CountRange(id*SliceWidth, (id+1)*SliceWidth) f.cache.BulkAdd(id, n) } f.cache.Invalidate() diff --git a/roaring/roaring.go b/roaring/roaring.go index 67799631f..14818f683 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -197,9 +197,14 @@ func (b *Bitmap) Count() (n uint64) { // CountRange returns the number of bits set between [start, end). func (b *Bitmap) CountRange(start, end uint64) (n uint64) { + if len(b.keys) == 0 { + return + } + skey := highbits(start) + ekey := highbits(end) - i := search64(b.keys, highbits(start)) - j := search64(b.keys, highbits(end)) + i := search64(b.keys, skey) + j := search64(b.keys, ekey) // If range is entirely in one container then just count that range. if i >= 0 && i == j { @@ -208,7 +213,13 @@ func (b *Bitmap) CountRange(start, end uint64) (n uint64) { // Count first partial container. if i < 0 { - i = -i + // start is before container, so we should start counting + // at first container that has value + if skey < b.keys[0] { + i = -1 + } else { + i = -i + } } else { n += uint64(b.containers[i].countRange(int(lowbits(start)), maxContainerVal+1)) } diff --git a/roaring/roaring_test.go b/roaring/roaring_test.go index 8de92eb4e..c721f3804 100644 --- a/roaring/roaring_test.go +++ b/roaring/roaring_test.go @@ -192,6 +192,29 @@ func TestBitmap_Max(t *testing.T) { } } +// Ensure CountRange is correct even if rangekey is prior to inital container. +func TestBitmap_BitmapCountRangeEdgeCase(t *testing.T) { + s := uint64(2009 * 1048576) + e := uint64(2010 * 1048576) + + start := s + (39314024 % 1048576) + bm0 := roaring.NewBitmap() + for i := uint64(0); i < 65536; i++ { + if (i+1)%4096 == 0 { + start += 16384 + } else { + start += 2 + } + bm0.Add(start) + } + a := bm0.Count() + r := bm0.CountRange(s, e) + if a != r { + t.Fatalf("Counts != CountRange %v %v", a, r) + + } +} + func TestBitmap_BitmapCountRange(t *testing.T) { bm0 := roaring.NewBitmap(0, 2683177) for i := uint64(628); i < 2683301; i++ {