diff --git a/executor.go b/executor.go index 6e46ad5b6..32d1b71eb 100644 --- a/executor.go +++ b/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) } diff --git a/fragment.go b/fragment.go index 6b846fe18..97dcf1be6 100644 --- a/fragment.go +++ b/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 diff --git a/pql/ast.go b/pql/ast.go index 39bf372e4..3480f4def 100644 --- a/pql/ast.go +++ b/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), diff --git a/roaring/filter.go b/roaring/filter.go index d7b91c892..9c79cc33f 100644 --- a/roaring/filter.go +++ b/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. diff --git a/roaring/filter_internal_test.go b/roaring/filter_internal_test.go index 6265ea0ae..3bffb4ef1 100644 --- a/roaring/filter_internal_test.go +++ b/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()