diff --git a/rbf/tx.go b/rbf/tx.go index d69e2c298..dff63bd1b 100644 --- a/rbf/tx.go +++ b/rbf/tx.go @@ -1300,6 +1300,7 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) { skey := highbits(start) ekey := highbits(end) + ebits := int32(lowbits(end)) csr, err := tx.cursor(name) if err == ErrBitmapNotFound { @@ -1339,7 +1340,7 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) { // If range is entirely in one container then just count that range. if skey == ekey { - return uint64(c.countRange(int32(lowbits(start)), int32(lowbits(end)))), nil + return uint64(c.countRange(int32(lowbits(start)), ebits)), nil } // INVAR: skey < ekey @@ -1356,8 +1357,8 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) { n += uint64(c.BitN) continue } - if k == ekey { - n += uint64(c.countRange(0, int32(lowbits(end)))) + if k == ekey && ebits > 0 { + n += uint64(c.countRange(0, ebits)) break } } diff --git a/roaring/roaring.go b/roaring/roaring.go index f0dffd560..2705a6ba3 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -1596,12 +1596,13 @@ func (b *Bitmap) removeEmptyContainers() { } } } -func (b *Bitmap) countEmptyContainers() int { + +func (b *Bitmap) countNonEmptyContainers() int { result := 0 citer, _ := b.Containers.Iterator(0) for citer.Next() { _, c := citer.Value() - if c.N() == 0 { + if c.N() > 0 { result++ } } @@ -1663,8 +1664,7 @@ func (b *Bitmap) writeToUnoptimized(w io.Writer) (n int64, err error) { // Remove empty containers before persisting. //b.removeEmptyContainers() - containerCount := b.Containers.Size() - b.countEmptyContainers() - headerSize := headerBaseSize + containerCount := b.countNonEmptyContainers() byte2 := make([]byte, 2) byte4 := make([]byte, 4) byte8 := make([]byte, 8) @@ -1695,12 +1695,17 @@ func (b *Bitmap) writeToUnoptimized(w io.Writer) (n int64, err error) { ew.WriteUint16(byte2, uint16(c.typ())) ew.WriteUint16(byte2, uint16(c.N()-1)) } - + } + if ew.n != headerBaseSize+(containerCount*12) { + return int64(ew.n), fmt.Errorf("after writing %d headers, wrote %d bytes, expected %d", + containerCount, ew.n, headerBaseSize+(containerCount*12)) } // Offset header section: write the offset for each container block. - // 4 bytes per container. - offset := uint32(headerSize + (containerCount * (8 + 2 + 2 + 4))) + // 4 bytes per container. The actual data offsets will then be 4 bytes + // further in per non-empty container, and now that we've scanned the + // containers, we know how many we really have. + offset := uint32(ew.n + (containerCount * 4)) citer, _ = b.Containers.Iterator(0) for citer.Next() { _, c := citer.Value() @@ -1708,13 +1713,18 @@ func (b *Bitmap) writeToUnoptimized(w io.Writer) (n int64, err error) { ew.WriteUint32(byte4, offset) offset += uint32(c.size()) } - } if ew.err != nil { return int64(ew.n), ew.err } + if ew.n != headerBaseSize+(containerCount*16) { + return int64(ew.n), fmt.Errorf("after writing %d headers+offsets, wrote %d bytes, expected %d", + containerCount, ew.n, headerBaseSize+(containerCount*16)) + } - n = int64(headerSize + (containerCount * (8 + 2 + 2 + 4))) + // We could compute the expected value, but a SliceContainers can contain + // a nil *Container, reported by Size() but not returned by the iterator. + n = int64(ew.n) // Container storage section: write each container block. citer, _ = b.Containers.Iterator(0) diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go index 0d1dbe7d8..9902c4373 100644 --- a/roaring/roaring_internal_test.go +++ b/roaring/roaring_internal_test.go @@ -2642,12 +2642,12 @@ func TestBitmap_RemoveEmptyContainers(t *testing.T) { bm1 := NewFileBitmap(1<<16, 2<<16, 3<<16) bm2 := NewFileBitmap(1<<16, 2<<16+1, 3<<16) bm3 := bm1.Intersect(bm2) - if bm3.countEmptyContainers() != 1 { + if bm3.countNonEmptyContainers() != 2 { t.Fatalf("Should be 1 empty container ") } bm3.removeEmptyContainers() - if bm3.countEmptyContainers() != 0 { + if bm3.countNonEmptyContainers() != bm3.Containers.Size() { t.Fatalf("Should be no empty containers ") } } @@ -2665,7 +2665,7 @@ func TestBitmap_BitmapWriteToWithEmpty(t *testing.T) { if err := bm0.UnmarshalBinary(buf.Bytes()); err != nil { t.Fatalf("unmarshalling: %v", err) } - if bm0.countEmptyContainers() != 0 { + if bm0.countNonEmptyContainers() != bm0.Containers.Size() { t.Fatalf("Should be no empty containers ") } if bm0.Count() != bm1.Count() { diff --git a/tx_internal_test.go b/tx_internal_test.go new file mode 100644 index 000000000..3cafba24a --- /dev/null +++ b/tx_internal_test.go @@ -0,0 +1,128 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pilosa + +import ( + "bytes" + "sync" + "testing" + + "github.com/pilosa/pilosa/v2/roaring" +) + +const countRangeMaxN = 8192 + +var countRangeSampleData []byte +var prepareCountRangeSampleData sync.Once + +// The sample data for the counter is just a series of containers, +// each with cardinality equal to its container key. +func requireCountRangeSampleData(tb testing.TB) (*fragment, Tx) { + prepareCountRangeSampleData.Do(func() { + var arraySample [4096]uint16 + // This horrible hack relies on a quirk of roaring's internals: It'll + // copy the bitmap if its length isn't exactly 1024. This lets us + // request that each container get its own copy of the bitmap. + var bitmapSample [1025]uint64 + for i := range arraySample { + arraySample[i] = uint16(i) + } + for i := 0; i < 4096/64; i++ { + bitmapSample[i] = ^uint64(0) + } + bm := roaring.NewSliceBitmap() + for n := 0; n < 4096 && n < countRangeMaxN; n++ { + c := roaring.NewContainerArray(arraySample[:n]) + bm.Put(uint64(n), c) + } + for n := 4096; n < countRangeMaxN; n++ { + c := roaring.NewContainerBitmapN(bitmapSample[:], int32(n)) + bm.Put(uint64(n), c) + bitmapSample[n/64] |= 1 << (n % 64) + } + var asBytes bytes.Buffer + n, err := bm.WriteTo(&asBytes) + if err != nil { + tb.Fatalf("writing bitmap: %v", err) + } + countRangeSampleData = asBytes.Bytes() + tb.Logf("creating bitmap: %d containers, %d bytes of data", countRangeMaxN, n) + }) + f, idx, tx := mustOpenFragment(tb, "i", "f", viewStandard, 0, "") + // Properly close this transaction, but not the next one we create that the + // caller will be responsible for. The deferred callback will + // be a nop if the Commit happened. + defer tx.Rollback() + err := f.importRoaringT(tx, countRangeSampleData, false) + if err != nil { + tb.Fatalf("importing sample data: %v", err) + } + err = tx.Commit() + if err != nil { + tb.Fatalf("committing sample data: %v", err) + } + tx = idx.holder.txf.NewTx(Txo{Write: false, Index: idx, Fragment: f, Shard: 0}) + return f, tx +} + +func TestTx_CountRange(t *testing.T) { + f, tx := requireCountRangeSampleData(t) + defer f.Clean(t) + defer tx.Rollback() + + expected := uint64(0) + j := uint64(0) + for i := uint64(0); i < countRangeMaxN; i += 7 { + if i%4 == 3 { + expected -= (j * 7) + 21 + j += 7 + } + got, err := tx.CountRange("i", "f", viewStandard, 0, uint64(j)<<16, uint64(i)<<16) + if err != nil { + t.Fatalf("counting range: %v", err) + } + if got != expected { + t.Fatalf("counting from container %d to %d, expected %d, got %d", + j, i, expected, got) + } + expected += (i * 7) + 21 + } +} + +func BenchmarkTx_CountRange(b *testing.B) { + f, tx := requireCountRangeSampleData(b) + defer f.Clean(b) + defer tx.Rollback() + + for k := 0; k < b.N; k++ { + expected := uint64(0) + j := uint64(0) + for i := uint64(0); i < countRangeMaxN; i += 7 { + if i%4 == 3 { + expected -= (j * 7) + 21 + j += 7 + } + got, err := tx.CountRange("i", "f", viewStandard, 0, uint64(j)<<16, uint64(i)<<16) + if err != nil { + b.Fatalf("counting range: %v", err) + } + if got != expected { + b.Fatalf("counting from container %d to %d, expected %d, got %d", + j, i, expected, got) + } + expected += (i * 7) + 21 + } + } +}