Merge pull request #1774 from richardartoul/ra/fix-union-in-place-bug

Fix bug in UnionInPlace function and add property test

roaring/roaring.go

@@ -617,7 +617,7 @@ func (b *Bitmap) unionInPlace(others ...*Bitmap) {
 					n:             maxContainerVal + 1,
 				}
 				target.Containers.Put(iKey, container)
-				bitmapIters[i:].markItersWithCurrentKeyAsHandled(iKey)
+				bitmapIters.markItersWithCurrentKeyAsHandled(i, iKey)
 				continue
 			}
 
@@ -632,8 +632,9 @@ func (b *Bitmap) unionInPlace(others ...*Bitmap) {
 			// 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 {
+			// otherwise we have to allocate a new one or convert the existing container
+			// into a bitmap container.
+			if container == nil {
 				buf := make([]uint64, bitmapN)
 				ob := buf[:bitmapN]
 				container = &Container{
@@ -641,6 +642,10 @@ func (b *Bitmap) unionInPlace(others ...*Bitmap) {
 					n:             0,
 					containerType: containerBitmap,
 				}
+			} else if container.isArray() {
+				container.arrayToBitmap()
+			} else if container.isRun() {
+				container.runToBitmap()
 			}
 
 			// Once we've acquired a bitmap container (either by reusing the existing one
@@ -648,7 +653,8 @@ func (b *Bitmap) unionInPlace(others ...*Bitmap) {
 			// 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:] {
+			for j := i; j < len(bitmapIters); j++ {
+				jIter := bitmapIters[j]
 				jKey, jContainer := jIter.iter.Value()
 
 				if iKey == jKey {
@@ -4002,8 +4008,9 @@ func (w handledIters) next() bool {
 	return hasNext
 }
 
-func (w handledIters) markItersWithCurrentKeyAsHandled(key uint64) {
-	for i, wrapped := range w {
+func (w handledIters) markItersWithCurrentKeyAsHandled(startIdx int, key uint64) {
+	for i := startIdx; i < len(w); i++ {
+		wrapped := w[i]
 		currKey, _ := wrapped.iter.Value()
 		if currKey == key {
 			w[i].handled = true

roaring/roaring_test.go

@@ -23,6 +23,7 @@ import (
 	"sort"
 	"testing"
 	"testing/quick"
+	"time"
 
 	"github.com/pilosa/pilosa"
 	"github.com/pilosa/pilosa/roaring"
@@ -436,6 +437,94 @@ func TestBitmap_UnionInPlace1(t *testing.T) {
 	}
 }
 
+// TestBitmap_UnionInPlaceProp is a manual property test that randomly generates
+// a number of different bitmaps with random vals and unions them together. It
+// then compares the result against a reference implementation (golang map) to
+// ensure that all the unions were handled correctly.
+func TestBitmap_UnionInPlaceProp(t *testing.T) {
+	var (
+		seed               = time.Now().UnixNano()
+		source             = rand.NewSource(seed)
+		rng                = rand.New(source)
+		numTests           = 100
+		maxNumIntsPerBatch = 100
+		maxNumBatches      = 100
+		maxRangePercent    = 2
+		// Need to limit the range of possible numbers that we generate
+		// otherwise two randomly generated numbers landing in the same
+		// container would be extremely unlikely, leaving container merging
+		// behavior untested.
+		maxUint64Val = 1000000
+	)
+
+	for i := 0; i < numTests; i++ {
+		var (
+			// We will use sets as the "reference" implementation.
+			sets    = []map[uint64]struct{}{}
+			bitmaps = []*roaring.Bitmap{}
+		)
+
+		// Ensure there are at least two batches.
+		numBatches := rng.Intn(maxNumBatches) + 2
+		for j := 0; j < numBatches; j++ {
+			// For each "batch" create the equivalent set and bitmap.
+			var (
+				set    = map[uint64]struct{}{}
+				bitmap = roaring.NewBitmap()
+			)
+
+			if rng.Intn(100) <= maxRangePercent {
+				// Generate max range RLE containers with a configurable
+				// probability to ensure that code-path is exercised.
+				start := rng.Intn((maxUint64Val))
+				// Add a continuous sequence of numbers that is 2x as long as the maximum
+				// size of a container to ensure we generate a maxRange container.
+				for x := start; x < (start + 2*(0xffff+1)); x++ {
+					set[uint64(x)] = struct{}{}
+					bitmap.Add(uint64(x))
+				}
+			}
+
+			// Generate and add a bunch of random values.
+			numIntsPerBatch := rng.Intn(maxNumIntsPerBatch)
+			for x := 0; x < numIntsPerBatch; x++ {
+				num := uint64(rng.Intn(maxUint64Val))
+				set[num] = struct{}{}
+				bitmap.Add(num)
+			}
+
+			sets = append(sets, set)
+			bitmaps = append(bitmaps, bitmap)
+		}
+
+		// "Union" all the sets into the first one.
+		set0 := sets[0]
+		for _, set := range sets[1:] {
+			for val := range set {
+				set0[val] = struct{}{}
+			}
+		}
+
+		// Union all the bitmaps into the first one.
+		bitmap0 := bitmaps[0]
+		bitmap0.UnionInPlace(bitmaps[1:]...)
+
+		// Ensure the unioned set and bitmap have the same cardinality.
+		if len(set0) != int(bitmap0.Count()) {
+			t.Fatalf("cardinality of set is: %d, but bitmap is: %d, failed with seed: %d",
+				len(set0), bitmap0.Count(), seed)
+		}
+
+		// Ensure the unioned set and bitmap have the exact same values.
+		for val := range set0 {
+			if !bitmap0.Contains(val) {
+				t.Fatalf("set contained %d, but bitmap did not, failed with seed: %d",
+					val, seed)
+			}
+		}
+	}
+}
+
 func TestBitmap_Intersection_Empty(t *testing.T) {
 	bm0 := roaring.NewFileBitmap(0, 2683177)
 	bm1 := roaring.NewFileBitmap()