From ff153459ee41552234e5c2a010494aadd4eed5ba Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Wed, 17 Jul 2019 23:09:38 -0500
Subject: [PATCH] union many things at once to cut down allocations

---
 executor.go |  8 ++++---
 row.go      | 63 +++++++++++++++++++++++++++++++++++++++--------------
 2 files changed, 52 insertions(+), 19 deletions(-)

diff --git a/executor.go b/executor.go
index 2110a92c6..57a206d90 100644
--- a/executor.go
+++ b/executor.go
@@ -1512,14 +1512,16 @@ func (e *executor) executeRowShard(ctx context.Context, index string, c *pql.Cal
 	}
 
 	// Union bitmaps across all time-based views.
-	row := &Row{}
-	for _, view := range viewsByTimeRange(viewStandard, fromTime, toTime, q) {
+	views := viewsByTimeRange(viewStandard, fromTime, toTime, q)
+	rows := make([]*Row, 0, len(views))
+	for _, view := range views {
 		f := e.Holder.fragment(index, fieldName, view, shard)
 		if f == nil {
 			continue
 		}
-		row = row.Union(f.row(rowID))
+		rows = append(rows, f.row(rowID))
 	}
+	row := rows[0].Union(rows[1:]...)
 	f.Stats.Count("range", 1, 1.0)
 	return row, nil
 
diff --git a/row.go b/row.go
index 79c918311..09a3390dc 100644
--- a/row.go
+++ b/row.go
@@ -150,21 +150,48 @@ func (r *Row) Xor(other *Row) *Row {
 }
 
 // Union returns the bitwise union of r and other.
-func (r *Row) Union(other *Row) *Row {
-	var segments []rowSegment
-	itr := newMergeSegmentIterator(r.segments, other.segments)
-	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
-		if s1 == nil {
-			segments = append(segments, *s0)
-			continue
-		} else if s0 == nil {
-			segments = append(segments, *s1)
-			continue
-		}
-		segments = append(segments, *s0.Union(s1))
+func (r *Row) Union(others ...*Row) *Row {
+	segments := make([][]rowSegment, 0, len(others)+1)
+	if len(r.segments) > 0 {
+		segments = append(segments, r.segments)
 	}
-
-	return &Row{segments: segments}
+	nextSegs := make([][]rowSegment, 0, len(others)+1)
+	toProcess := make([]*rowSegment, 0, len(others)+1)
+	var output []rowSegment
+	for _, other := range others {
+		if len(other.segments) > 0 {
+			segments = append(segments, other.segments)
+		}
+	}
+	for len(segments) > 0 {
+		shard := segments[0][0].shard
+		for _, segs := range segments {
+			if segs[0].shard < shard {
+				shard = segs[0].shard
+			}
+		}
+		nextSegs = nextSegs[:0]
+		toProcess := toProcess[:0]
+		for _, segs := range segments {
+			if segs[0].shard == shard {
+				toProcess = append(toProcess, &segs[0])
+				segs = segs[1:]
+			}
+			if len(segs) > 0 {
+				nextSegs = append(nextSegs, segs)
+			}
+		}
+		// at this point, "toProcess" is a list of all the segments
+		// sharing the lowest ID, and nextSegs is a list of all the others.
+		// Swap the segment lists (so we don't have to reallocate it)
+		segments, nextSegs = nextSegs, segments
+		if len(toProcess) == 1 {
+			output = append(output, *toProcess[0])
+		} else {
+			output = append(output, *toProcess[0].Union(toProcess[1:]...))
+		}
+	}
+	return &Row{segments: output}
 }
 
 // Difference returns the diff of r and other.
@@ -350,8 +377,12 @@ func (s *rowSegment) Intersect(other *rowSegment) *rowSegment {
 }
 
 // Union returns the bitwise union of s and other.
-func (s *rowSegment) Union(other *rowSegment) *rowSegment {
-	data := s.data.Union(other.data)
+func (s *rowSegment) Union(others ...*rowSegment) *rowSegment {
+	datas := make([]*roaring.Bitmap, len(others))
+	for i, other := range others {
+		datas[i] = other.data
+	}
+	data := s.data.Union(datas...)
 	data.Freeze()
 
 	return &rowSegment{