Union run containers by run instead of by value

arraycontainer.go

@@ -428,12 +428,20 @@ func (ac *arrayContainer) iorBitmap(bc2 *bitmapContainer) container {
 	return bc1
 }
 
+// arrayRunInsertMaxRuns is the most runs iorRun16 inserts into the array one
+// at a time. Each insert moves the tail of the array, so past a few runs one
+// merge is usually cheaper.
+const arrayRunInsertMaxRuns = 4
+
 func (ac *arrayContainer) iorRun16(rc *runContainer16) container {
 	runCardinality := rc.getCardinality()
 	// heuristic for if the container should maybe be an
 	// array container.
 	if runCardinality < ac.getCardinality() &&
 		runCardinality+ac.getCardinality() < arrayDefaultMaxSize {
+		if len(rc.iv) > arrayRunInsertMaxRuns {
+			return ac.iorArray(rc.toArrayContainer())
+		}
 		var result container
 		result = ac
 		for _, run := range rc.iv {

real_data_benchmark_test.go

@@ -186,3 +186,37 @@ func BenchmarkRealDataFastOr(b *testing.B) {
 		return FastOr(bitmaps...).GetCardinality()
 	})
 }
+
+func BenchmarkRealDataOrInPlace(b *testing.B) {
+	benchmarkRealDataAggregate(b, func(bitmaps []*Bitmap) uint64 {
+		tot := uint64(0)
+		for i := 1; i < len(bitmaps); i++ {
+			x := bitmaps[i-1].Clone()
+			x.Or(bitmaps[i])
+			tot += x.GetCardinality()
+		}
+		return tot
+	})
+}
+
+func BenchmarkRealDataOrFold(b *testing.B) {
+	benchmarkRealDataAggregate(b, func(bitmaps []*Bitmap) uint64 {
+		x := bitmaps[0].Clone()
+		for _, y := range bitmaps[1:] {
+			x.Or(y)
+		}
+		return x.GetCardinality()
+	})
+}
+
+func BenchmarkRealDataAndInPlace(b *testing.B) {
+	benchmarkRealDataAggregate(b, func(bitmaps []*Bitmap) uint64 {
+		tot := uint64(0)
+		for i := 1; i < len(bitmaps); i++ {
+			x := bitmaps[i-1].Clone()
+			x.And(bitmaps[i])
+			tot += x.GetCardinality()
+		}
+		return tot
+	})
+}

runcontainer.go

@@ -417,12 +417,15 @@ func (rc *runContainer16) union(b *runContainer16) *runContainer16 {
 		} else {
 			// !mergedUsed
 			if !canMerge16(cura, curb) {
+				// copy every run that ends before the other side's next run
 				if cura.start < curb.start {
-					m = append(m, cura)
-					na++
+					next := rc.indexOfIntervalAtOrAfter(int(curb.start)-1, na+1)
+					m = append(m, rc.iv[na:next]...)
+					na = next
 				} else {
-					m = append(m, curb)
-					nb++
+					next := b.indexOfIntervalAtOrAfter(int(cura.start)-1, nb+1)
+					m = append(m, b.iv[nb:next]...)
+					nb = next
 				}
 			} else {
 				merged = mergeInterval16s(cura, curb)
@@ -595,6 +598,11 @@ func (rc *runContainer16) unionCardinality(b *runContainer16) uint {
 
 // indexOfIntervalAtOrAfter is a helper for union.
 func (rc *runContainer16) indexOfIntervalAtOrAfter(key int, startIndex int) int {
+	// union only asks for keys past rc.iv[startIndex-1], so a key at or
+	// before the end of rc.iv[startIndex] needs no search.
+	if startIndex < len(rc.iv) && int(rc.iv[startIndex].last()) >= key {
+		return startIndex
+	}
 	w, already, _ := rc.searchRange(key, startIndex, 0)
 	if already {
 		return w
@@ -1418,6 +1426,9 @@ func intersectWithLeftover16(astart, alast, bstart, blast int) (isOverlap, isLef
 }
 
 func (rc *runContainer16) findNextIntervalThatIntersectsStartingFrom(startIndex int, key int) (index int, done bool) {
+	if startIndex+1 < len(rc.iv) && int(rc.iv[startIndex+1].start) > key {
+		return startIndex, false
+	}
 	w, _, _ := rc.searchRange(key, startIndex, 0)
 	// rc.search always returns w < len(rc.iv)
 	if w < startIndex {
@@ -2267,16 +2278,53 @@ func (rc *runContainer16) ior(a container) container {
 	panic("unsupported container type")
 }
 
+// runInsertMaxRuns is the most runs inplaceUnion inserts one at a time. Each
+// insert shifts the runs after it, so past a few a merge is usually cheaper.
+const runInsertMaxRuns = 16
+
 func (rc *runContainer16) inplaceUnion(rc2 *runContainer16) container {
-	for _, p := range rc2.iv {
-		last := int(p.last())
-		for i := int(p.start); i <= last; i++ {
-			rc.Add(uint16(i))
+	if len(rc2.iv) <= runInsertMaxRuns {
+		for _, p := range rc2.iv {
+			rc.addInterval(p)
 		}
+	} else {
+		*rc = *rc.union(rc2)
 	}
 	return rc.toEfficientContainer()
 }
 
+// addInterval adds the values of iv in place, merging the runs it touches.
+func (rc *runContainer16) addInterval(iv interval16) {
+	start, last := int(iv.start), int(iv.last())
+	// runs i..j-1 overlap or touch iv
+	i, j := 0, len(rc.iv)
+	for i < j {
+		h := int(uint(i+j) >> 1)
+		if int(rc.iv[h].last())+1 < start {
+			i = h + 1
+		} else {
+			j = h
+		}
+	}
+	j = len(rc.iv)
+	for lo := i; lo < j; {
+		h := int(uint(lo+j) >> 1)
+		if int(rc.iv[h].start) <= last+1 {
+			lo = h + 1
+		} else {
+			j = h
+		}
+	}
+	if i == j {
+		rc.iv = slices.Insert(rc.iv, i, iv)
+		return
+	}
+	start = min(start, int(rc.iv[i].start))
+	last = max(last, int(rc.iv[j-1].last()))
+	rc.iv[i] = newInterval16Range(uint16(start), uint16(last))
+	rc.iv = append(rc.iv[:i+1], rc.iv[j:]...)
+}
+
 // Such code should not be used as it will not preserve the container invariants:
 //func (rc *runContainer16) iorBitmapContainer(bc *bitmapContainer) container {
 //	it := bc.getShortIterator()
@@ -2717,9 +2765,13 @@ func (rc *runContainer16) toEfficientContainerFromCardinality(card int) containe
 }
 
 func (rc *runContainer16) toArrayContainer() *arrayContainer {
-	ac := newArrayContainer()
-	for i := range rc.iv {
-		ac.iaddRange(int(rc.iv[i].start), int(rc.iv[i].last())+1)
+	ac := newArrayContainerSize(rc.getCardinality())
+	pos := 0
+	for _, iv := range rc.iv {
+		for v := int(iv.start); v <= int(iv.last()); v++ {
+			ac.content[pos] = uint16(v)
+			pos++
+		}
 	}
 	return ac
 }

runcontainer_union_test.go

@@ -0,0 +1,161 @@
+package roaring
+
+import (
+	"fmt"
+	"math/rand"
+	"slices"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// unionValues is the reference: the sorted union of the values in a and b.
+func unionValues(a, b []uint16) []uint16 {
+	out := slices.Concat(a, b)
+	slices.Sort(out)
+	return slices.Compact(out)
+}
+
+func containerValues(c container) []uint16 {
+	var out []uint16
+	for it := c.getShortIterator(); it.hasNext(); {
+		out = append(out, it.next())
+	}
+	return out
+}
+
+// randomRuns returns a run container of up to maxRuns runs starting at or
+// after lo, each 1..maxLen long, separated by gaps of 1..maxGap.
+func randomRuns(r *rand.Rand, lo, maxRuns, maxLen, maxGap int) *runContainer16 {
+	var vals []uint16
+	v := lo
+	for k := 0; k < maxRuns && v < 1<<16; k++ {
+		for end := min(v+1+r.Intn(maxLen), 1<<16); v < end; v++ {
+			vals = append(vals, uint16(v))
+		}
+		v += 1 + r.Intn(maxGap)
+	}
+	return newRunContainer16FromVals(true, vals...)
+}
+
+func TestRunContainerUnionExhaustive(t *testing.T) {
+	// every pair of subsets of an 8-value window, at both ends of the range
+	for _, base := range []int{0, 1<<16 - 8} {
+		for x := 0; x < 256; x++ {
+			for y := 0; y < 256; y++ {
+				var av, bv []uint16
+				for i := 0; i < 8; i++ {
+					if x&(1<<i) != 0 {
+						av = append(av, uint16(base+i))
+					}
+					if y&(1<<i) != 0 {
+						bv = append(bv, uint16(base+i))
+					}
+				}
+				a, b := newRunContainer16FromVals(true, av...), newRunContainer16FromVals(true, bv...)
+				want := unionValues(av, bv)
+				require.Equal(t, want, containerValues(a.union(b)))
+				var both []uint16
+				for _, v := range av {
+					if slices.Contains(bv, v) {
+						both = append(both, v)
+					}
+				}
+				require.Equal(t, both, containerValues(a.intersect(b)))
+				got := a.Clone().inplaceUnion(b)
+				require.Equal(t, want, containerValues(got))
+				if len(want) > 0 {
+					require.NoError(t, got.validate())
+				}
+			}
+		}
+	}
+}
+
+func TestRunContainerInplaceUnionRandom(t *testing.T) {
+	r := rand.New(rand.NewSource(1))
+	for trial := 0; trial < 500; trial++ {
+		maxLen, maxGap := 1+r.Intn(100), 1+r.Intn(100)
+		a := randomRuns(r, r.Intn(64), 1+r.Intn(3000), maxLen, maxGap)
+		// small operands straddle runInsertMaxRuns; every other trial uses a large one
+		b := randomRuns(r, r.Intn(1<<16), 1+r.Intn(2*runInsertMaxRuns), maxLen, maxGap)
+		if trial%2 == 1 {
+			b = randomRuns(r, r.Intn(64), 1+r.Intn(3000), maxLen, maxGap)
+		}
+		av, bv := containerValues(a), containerValues(b)
+
+		got := a.inplaceUnion(b)
+		require.Equal(t, unionValues(av, bv), containerValues(got))
+		require.NoError(t, got.validate())
+		require.Equal(t, bv, containerValues(b), "operand modified")
+	}
+	// self-union
+	a := randomRuns(r, 0, 10, 50, 50)
+	want := containerValues(a)
+	require.Equal(t, want, containerValues(a.inplaceUnion(a)))
+}
+
+func TestArrayContainerIorRun16Random(t *testing.T) {
+	r := rand.New(rand.NewSource(2))
+	for trial := 0; trial < 2000; trial++ {
+		ac := newArrayContainer()
+		for n := 1 + r.Intn(arrayDefaultMaxSize); n > 0; n-- {
+			ac.iadd(uint16(r.Intn(1 << 16)))
+		}
+		rc := randomRuns(r, r.Intn(1<<16), 1+r.Intn(4*arrayRunInsertMaxRuns), 1+r.Intn(20), 1+r.Intn(2000))
+		acv, rcv := containerValues(ac), containerValues(rc)
+
+		got := ac.iorRun16(rc)
+		require.Equal(t, unionValues(acv, rcv), containerValues(got))
+		require.Equal(t, rcv, containerValues(rc), "operand modified")
+		if len(rcv) < len(acv) && len(rcv)+len(acv) < arrayDefaultMaxSize {
+			require.Same(t, ac, got)
+		}
+	}
+}
+
+func TestRunContainerToArrayContainer(t *testing.T) {
+	r := rand.New(rand.NewSource(3))
+	for trial := 0; trial < 500; trial++ {
+		rc := randomRuns(r, r.Intn(1000), 1+r.Intn(200), 1+r.Intn(40), 1+r.Intn(300))
+		require.Equal(t, containerValues(rc), rc.toArrayContainer().content)
+	}
+}
+
+func BenchmarkRunContainerInplaceUnion(b *testing.B) {
+	receiver := newRunContainer16()
+	for v := 0; v < 1<<16; v += 64 {
+		receiver.iv = append(receiver.iv, newInterval16Range(uint16(v), uint16(v+31)))
+	}
+	for _, runs := range []int{4, 1024} {
+		operand := newRunContainer16()
+		for i := 0; i < runs; i++ {
+			v := i * (1 << 16) / runs
+			operand.iv = append(operand.iv, newInterval16Range(uint16(v+40), uint16(v+47)))
+		}
+		b.Run(fmt.Sprintf("runs=%d", runs), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				receiver.Clone().inplaceUnion(operand)
+			}
+		})
+	}
+}
+
+func BenchmarkArrayContainerIorRun16(b *testing.B) {
+	ac := newArrayContainer()
+	for v := 0; v < 1<<16; v += 32 {
+		ac.content = append(ac.content, uint16(v))
+	}
+	for _, runs := range []int{4, 100} {
+		rc := newRunContainer16()
+		for i := 0; i < runs; i++ {
+			v := i*(1<<16)/runs + 7
+			rc.iv = append(rc.iv, newInterval16Range(uint16(v), uint16(v+2)))
+		}
+		b.Run(fmt.Sprintf("runs=%d", runs), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				ac.clone().(*arrayContainer).iorRun16(rc)
+			}
+		})
+	}
+}