Implement UnionRows as a BitmapFilter (#2116)

This is fairly experimental, but basically, we make a fragment-level
op which, given a []uint64, can produce a union of all those rows in
the fragment, with a single scan through the fragment and not needing
a ton of additional space to reify all the rows at once.

Now with the Repair calls happening in the Results assembly rather than
on the intermediate data.

executor.go

@@ -1490,6 +1490,8 @@ func (e *executor) executeBitmapCallShard(ctx context.Context, qcx *Qcx, index s
 		return e.executeIntersectShard(ctx, qcx, index, c, shard)
 	case "Union":
 		return e.executeUnionShard(ctx, qcx, index, c, shard)
+	case "InnerUnionRows":
+		return e.executeInnerUnionRowsShard(ctx, qcx, index, c, shard)
 	case "Xor":
 		return e.executeXorShard(ctx, qcx, index, c, shard)
 	case "Not":
@@ -4856,6 +4858,116 @@ func (e *executor) executeUnionShard(ctx context.Context, qcx *Qcx, index string
 	return rows[0].Union(rows[1:]...), nil
 }
 
+// executeInnerUnionRowsShard executes a special magical call which is actually
+// more like Row() than Union(), and takes a call plus a []uint64 of rows, and
+// generates the union of the rows in the []uint64.
+func (e *executor) executeInnerUnionRowsShard(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shard uint64) (out *Row, err0 error) {
+	span, _ := tracing.StartSpanFromContext(ctx, "Executor.executeInnerUnionRowsShard")
+	defer span.Finish()
+
+	// Fetch index.
+	idx := e.Holder.Index(index)
+	if idx == nil {
+		return nil, newNotFoundError(ErrIndexNotFound, index)
+	}
+
+	fieldName, ok, err := c.StringArg("_field")
+	if err != nil {
+		return nil, errors.Wrap(err, "finding field")
+	}
+	if !ok {
+		return nil, errors.New("InnerUnionRows requires _field")
+	}
+
+	f := idx.Field(fieldName)
+	if f == nil {
+		return nil, newNotFoundError(ErrFieldNotFound, fieldName)
+	}
+
+	// Parse "from" time, if set.
+	var fromTime time.Time
+	if v, ok := c.Args["from"]; ok {
+		if fromTime, err = parseTime(v); err != nil {
+			return nil, errors.Wrap(err, "parsing from time")
+		}
+	}
+
+	// Parse "to" time, if set.
+	var toTime time.Time
+	if v, ok := c.Args["to"]; ok {
+		if toTime, err = parseTime(v); err != nil {
+			return nil, errors.Wrap(err, "parsing to time")
+		}
+	}
+
+	rowIDs, rowOK, err := c.UintSliceArg("rows")
+	if err != nil {
+		return nil, fmt.Errorf("extracting rows argument: %v", err)
+	}
+	if !rowOK {
+		return nil, fmt.Errorf("InnerUnionRows() must specify rows")
+	}
+
+	// Simply return row if times are not set.
+	timeNotSet := fromTime.IsZero() && toTime.IsZero()
+	if timeNotSet {
+		frag := e.Holder.fragment(index, fieldName, viewStandard, shard)
+		if frag == nil {
+			return NewRow(), nil
+		}
+
+		tx, finisher, err := qcx.GetTx(Txo{Write: !writable, Fragment: frag, Index: idx, Shard: shard})
+		if err != nil {
+			return nil, err
+		}
+		defer finisher(&err0)
+		row, err := frag.unionRows(ctx, tx, rowIDs)
+		if qcx.write && err == nil {
+			row = row.Clone()
+		}
+		return row, err
+	}
+
+	views, err := f.viewsByTimeRange(fromTime, toTime)
+	if err != nil {
+		return nil, err
+	}
+
+	// Union bitmaps across all time-based views.
+	rows := make([]*Row, 0, len(views))
+	tx, finisher, err := qcx.GetTx(Txo{Write: !writable, Index: idx, Shard: shard})
+	defer finisher(&err0)
+	for _, view := range views {
+		f := e.Holder.fragment(index, fieldName, view, shard)
+		if f == nil {
+			continue
+		}
+		if err != nil {
+			return nil, err
+		}
+
+		row, err := f.unionRows(ctx, tx, rowIDs)
+		if err != nil {
+			return nil, err
+		}
+		rows = append(rows, row)
+	}
+	if len(rows) == 0 {
+		return &Row{}, nil
+	} else if len(rows) == 1 {
+		if qcx.write {
+			return rows[0].Clone(), nil
+		}
+		return rows[0], nil
+	}
+	row := rows[0].Union(rows[1:]...)
+	if qcx.write {
+		row = row.Clone()
+	}
+	return row, nil
+
+}
+
 // executeXorShard executes a xor() call for a local shard.
 func (e *executor) executeXorShard(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shard uint64) (_ *Row, err error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeXorShard")
@@ -4981,6 +5093,8 @@ func (e *executor) executeUnionRows(ctx context.Context, qcx *Qcx, index string,
 		switch rowsResult := rowsResult.(type) {
 		case *PairsField:
 			// Translate pairs into rows calls.
+			// TODO: This should probably also be adjusted to use InnerUnionRows,
+			// but we can't do that for the string key case.
 			for _, p := range rowsResult.Pairs {
 				var val interface{}
 				switch {
@@ -4997,15 +5111,14 @@ func (e *executor) executeUnionRows(ctx context.Context, qcx *Qcx, index string,
 				})
 			}
 		case RowIDs:
-			// Translate Row IDs into Row calls.
-			for _, id := range rowsResult {
-				resultRows = append(resultRows, &pql.Call{
-					Name: "Row",
-					Args: map[string]interface{}{
-						child.Args["_field"].(string): id,
-					},
-				})
-			}
+			// Make a single InnerUnionRows from this
+			resultRows = append(resultRows, &pql.Call{
+				Name: "InnerUnionRows",
+				Args: map[string]interface{}{
+					"_field": child.Args["_field"],
+					"rows":   []uint64(rowsResult),
+				},
+			})
 		default:
 			return nil, errors.Errorf("unexpected Rows type %T", rowsResult)
 		}

fragment.go

@@ -2820,6 +2820,32 @@ func (f *fragment) unprotectedRows(ctx context.Context, tx Tx, start uint64, fil
 	}
 }
 
+// unionRows yields the union of the given rows in this fragment
+func (f *fragment) unionRows(ctx context.Context, tx Tx, rows []uint64) (*Row, error) {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+	return f.unprotectedUnionRows(ctx, tx, rows)
+}
+
+// unprotectedRows calls rows without grabbing the mutex.
+func (f *fragment) unprotectedUnionRows(ctx context.Context, tx Tx, rows []uint64) (*Row, error) {
+	filter := roaring.NewBitmapRowsUnion(rows)
+	err := tx.ApplyFilter(f.index(), f.field(), f.view(), f.shard, 0, filter)
+	if err != nil {
+		return nil, err
+	} else {
+		row := &Row{
+			segments: []rowSegment{{
+				data:     filter.Results(f.shard),
+				shard:    f.shard,
+				writable: true,
+			}},
+		}
+		row.invalidateCount()
+		return row, nil
+	}
+}
+
 // blockToRoaringData converts a fragment block into a roaring.Bitmap
 // which represents a portion of the data within a single shard.
 // TODO: it seems like we should be able to get the

pql/ast.go

@@ -474,6 +474,16 @@ var callInfoByFunc = map[string]callInfo{
 			"in":       nil,
 		},
 	},
+	"InnerUnionRows": {
+		allowUnknown: false,
+		prototypes: map[string]interface{}{
+			"_field": stringOrVariable,
+			"field":  stringOrVariable,
+			"from":   nil,
+			"to":     nil,
+			"rows":   nil,
+		},
+	},
 	"Shift": {allowUnknown: false,
 		prototypes: map[string]interface{}{
 			"n": int64(0),

roaring/filter.go

@@ -290,6 +290,81 @@ func NewBitmapRowsFilter(rows []uint64) BitmapFilter {
 	return &BitmapRowsFilter{rows: rows, i: 0}
 }
 
+// BitmapRowsUnion is a BitmapFilter which produces a union of all the
+// rows listed in a []uint64.
+type BitmapRowsUnion struct {
+	c    []*Container
+	rows []uint64
+	i    int
+}
+
+func (f *BitmapRowsUnion) ConsiderKey(key FilterKey, n int32) FilterResult {
+	if f.i == -1 {
+		return key.Done()
+	}
+	if n == 0 {
+		return key.RejectOne()
+	}
+	row := uint64(key) >> rowExponent
+	for f.rows[f.i] < row {
+		f.i++
+		if f.i >= len(f.rows) {
+			f.i = -1
+			return key.Done()
+		}
+	}
+	if f.rows[f.i] > row {
+		return key.RejectUntilRow(f.rows[f.i])
+	}
+	// If we ran out of rows, we said we were done. If we're
+	// waiting for a later row, we said to reject until then.
+	// Therefore, we're on the current row, and need the data because
+	// we're going to union it.
+	return key.NeedData()
+}
+
+func (f *BitmapRowsUnion) ConsiderData(key FilterKey, data *Container) FilterResult {
+	idx := key & keyMask
+	f.c[idx] = f.c[idx].UnionInPlace(data)
+	// UnionInPlace with nil will reuse the container. We don't want to reuse
+	// the container, because ApplyFilter will overwrite it.
+	if f.c[idx] == data {
+		f.c[idx] = data.Clone()
+	}
+	return key.MatchOne()
+}
+
+// Yield the bitmap containing our results, adjusted for a particular shard
+// if necessary (because we expect the results to correspond to our shard
+// ID).
+func (f *BitmapRowsUnion) Results(shard uint64) *Bitmap {
+	shard <<= shardwidth.Exponent
+	b := NewSliceBitmap()
+	for i, c := range f.c {
+		// UnionInPlace might not have fixed count
+		c.Repair()
+		b.Containers.Put(uint64(i)+shard, c)
+	}
+	return b
+}
+
+// Reset the internal container buffer. You must use this before reusing a
+// filter.
+func (f *BitmapRowsUnion) Reset() {
+	for i := range f.c {
+		f.c[i] = nil
+	}
+}
+
+// NewBitmapRowsUnion yields a BitmapRowsUnion which can give you the union
+// of all containers matching a given row.
+func NewBitmapRowsUnion(rows []uint64) *BitmapRowsUnion {
+	if len(rows) == 0 {
+		return &BitmapRowsUnion{rows: rows, i: -1, c: make([]*Container, rowWidth)}
+	}
+	return &BitmapRowsUnion{rows: rows, i: 0, c: make([]*Container, rowWidth)}
+}
+
 // BitmapRowFilterBase is a generic form of a row-aware wrapper; it
 // handles making decisions about keys once you tell it a yesKey and noKey
 // that it should be using, and makes callbacks per row.

roaring/filter_internal_test.go

@@ -113,6 +113,20 @@ func TestRowsFilter(t *testing.T) {
 	compareSlices(t, "limit", expected[:1], rows)
 }
 
+func TestRowsUnion(t *testing.T) {
+	requireSampleData(t)
+	rowSet := []uint64{7, 11}
+	expected := []uint64{1<<16 + 1, 7<<16 + 7, 11<<16 + 11}
+	u := NewBitmapRowsUnion(rowSet)
+	iter, _ := filterSampleData.Containers.Iterator(0)
+	err := ApplyFilterToIterator(u, iter)
+	if err != nil {
+		t.Fatalf("unexpected filter error: %v", err)
+	}
+	out := u.Results(0)
+	compareSlices(t, "sevenEleven", out.Slice(), expected)
+}
+
 func TestBitmapFilter(t *testing.T) {
 	requireSampleData(t)
 	bm := NewBitmap()