Merge pull request #650 from jaddr2line/limit

Add a "Limit" query
This commit is contained in:
Nia 2020-08-12 14:04:12 -04:00 committed by GitHub
commit 50a2e3c5fa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 206 additions and 58 deletions

View file

@ -571,6 +571,38 @@ Not(Row(stargazer=1))
* columns are repositories that were not starred by user 1
#### Limit
**Spec:**
```
Limit(<ROW_CALL>, [limit=<UINT>], [offset=<UINT>])
```
**Description:**
Limit executes a `ROW_CALL` and returns a subset of the results.
If a limit of `n` is specified, then this query will return the first `n` results of the row call.
If an offset of `m` is specified, then this query will skip the first `m` results of the row call.
If both a limit and offset are specified, the offset is applied before the limit.
This can be used to implement pagination.
**Result Type:** object with attrs and columns
attrs will always be empty
**Examples:**
Find the second column that has a bit set in the given row.
```request
Limit(Row(stargazer=1), limit=1, offset=1)
```
```response
{"results":[{"attrs":{},"columns":[30]}]}
```
* columns are repositories that were starred by user 1
#### Count
**Spec:**

View file

@ -649,6 +649,43 @@ func (e *executor) preprocessQuery(ctx context.Context, tx Tx, index string, c *
Precomputed: precomputed,
}, nil
case "All":
_, hasLimit, err := c.UintArg("limit")
if err != nil {
return nil, err
}
_, hasOffset, err := c.UintArg("offset")
if err != nil {
return nil, err
}
if !hasLimit && !hasOffset {
return c, nil
}
// Rewrite the All() w/ limit to Limit(All()).
c.Children = []*pql.Call{
{
Name: "All",
},
}
c.Name = "Limit"
fallthrough
case "Limit":
if len(c.Children) != 1 {
return nil, errors.Errorf("expected 1 child of limit call but got %d", len(c.Children))
}
res, err := e.preprocessQuery(ctx, tx, index, c.Children[0], shards, opt)
if err != nil {
return nil, err
}
c.Children[0] = res
err = e.executeLimitCall(ctx, tx, index, c, shards, opt)
if err != nil {
return nil, err
}
return c, nil
default:
// Recurse through child calls.
out := make([]*pql.Call, len(c.Children))
@ -776,9 +813,6 @@ func (e *executor) executeCall(ctx context.Context, tx Tx, index string, c *pql.
case "FieldValue":
statFn()
return e.executeFieldValueCall(ctx, tx, index, c, shards, opt)
case "All":
statFn()
return e.executeAllCall(ctx, tx, index, c, shards, opt)
case "Precomputed":
return e.executePrecomputedCall(ctx, tx, index, c, shards, opt)
default:
@ -976,25 +1010,20 @@ func (e *executor) executeFieldValueCallShard(ctx context.Context, tx Tx, field
return other, nil
}
// executeAllCall executes an All() call.
func (e *executor) executeAllCall(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*Row, error) {
rslt := NewRow()
// executeLimitCall executes a Limit() call, **rewriting it to a precomputed call**.
func (e *executor) executeLimitCall(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) error {
bitmapCall := c.Children[0]
var limit uint64
var offset uint64
if lim, hasLimit, err := c.UintArg("limit"); err != nil {
return nil, errors.Wrap(err, "getting limit")
} else if hasLimit && lim > 0 {
limit = uint64(lim)
limit, hasLimit, err := c.UintArg("limit")
if err != nil {
return errors.Wrap(err, "getting limit")
}
if off, hasOffset, err := c.UintArg("offset"); err != nil {
return nil, errors.Wrap(err, "getting offset")
} else if hasOffset && off > 0 {
offset = uint64(off)
offset, _, err := c.UintArg("offset")
if err != nil {
return errors.Wrap(err, "getting offset")
}
if limit == 0 {
if !hasLimit {
limit = math.MaxUint64
}
@ -1005,12 +1034,34 @@ func (e *executor) executeAllCall(ctx context.Context, tx Tx, index string, c *p
// got tracks the number of records gotten to that point.
var got uint64
c.Precomputed = make(map[uint64]interface{})
for _, shard := range shards {
row, err := e.executeAllCallMapReduce(ctx, tx, index, c, shard, opt)
if err != nil {
return nil, errors.Wrap(err, "executing map reduce on shard")
// Execute calls in bulk on each remote node and merge.
mapFn := func(ctx context.Context, shard uint64) (interface{}, error) {
return e.executeBitmapCallShard(ctx, tx, index, bitmapCall, shard)
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
if err := ctx.Err(); err != nil {
return err
}
other, _ := prev.(*Row)
if other == nil {
other = NewRow()
}
other.Merge(v.(*Row))
return other
}
result, err := e.mapReduce(ctx, index, []uint64{shard}, c, opt, mapFn, reduceFn)
if err != nil {
return errors.Wrap(err, "limit map reduce")
}
row, _ := result.(*Row)
segCnt := row.Count()
// If this segment doesn't reach the offset, skip it.
@ -1023,14 +1074,14 @@ func (e *executor) executeAllCall(ctx context.Context, tx Tx, index string, c *p
// (or it has exactly enough).
if segCnt-skip <= limit-got {
if skip == 0 {
rslt.Merge(row)
c.Precomputed[shard] = row
} else {
cols := row.Columns()
partialRow := NewRow()
for _, bit := range cols[skip:] {
partialRow.SetBit(bit)
}
rslt.Merge(partialRow)
c.Precomputed[shard] = partialRow
}
got += segCnt - skip
// In the case where this segment exactly fulfills the limit, break.
@ -1047,42 +1098,12 @@ func (e *executor) executeAllCall(ctx context.Context, tx Tx, index string, c *p
for _, bit := range cols[skip : skip+limit-got] {
partialRow.SetBit(bit)
}
rslt.Merge(partialRow)
c.Precomputed[shard] = partialRow
break
}
return rslt, nil
}
// executeAllCallMapReduce executes a single shard of the All() call
// using the executor.mapReduce() method.
func (e *executor) executeAllCallMapReduce(ctx context.Context, tx Tx, index string, c *pql.Call, shard uint64, opt *execOptions) (*Row, error) {
// Execute calls in bulk on each remote node and merge.
mapFn := func(ctx context.Context, shard uint64) (interface{}, error) {
return e.executeAllCallShard(ctx, tx, index, c, shard)
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
if err := ctx.Err(); err != nil {
return err
}
other, _ := prev.(*Row)
if other == nil {
other = NewRow()
}
other.Merge(v.(*Row))
return other
}
result, err := e.mapReduce(ctx, index, []uint64{shard}, c, opt, mapFn, reduceFn)
if err != nil {
return nil, errors.Wrap(err, "map reduce")
}
row, _ := result.(*Row)
return row, nil
c.Name = "Precomputed"
return nil
}
// executeIncludesColumnCallShard
@ -1458,7 +1479,7 @@ func (e *executor) executeBitmapCallShard(ctx context.Context, tx Tx, index stri
return e.executeNotShard(ctx, tx, index, c, shard)
case "Shift":
return e.executeShiftShard(ctx, tx, index, c, shard)
case "All": // Allow a shard computation to use All() (note, limit/offset not applied)
case "All": // Allow a shard computation to use All()
return e.executeAllCallShard(ctx, tx, index, c, shard)
case "Distinct":
return nil, errors.New("Distinct shouldn't be hit as a bitmap call")

View file

@ -3650,6 +3650,94 @@ func TestExecutor_Execute_FieldValue(t *testing.T) {
}
}
// Ensure a Limit query can be executed.
func TestExecutor_Execute_Limit(t *testing.T) {
c := test.MustRunCluster(t, 1)
defer c.Close()
c.CreateField(t, "i", pilosa.IndexOptions{TrackExistence: true}, "f")
c.ImportBits(t, "i", "f", [][2]uint64{
{1, 0},
{1, 1},
{1, ShardWidth + 1},
})
columns := []uint64{0, 1, ShardWidth + 1}
// Test with only a limit specified.
t.Run("Limit", func(t *testing.T) {
for limit := 0; limit < 5; limit++ {
expect := columns
if limit < len(expect) {
expect = expect[:limit]
}
resp := c.Query(t, "i", fmt.Sprintf("Limit(All(), limit=%d)", limit))
if len(resp.Results) != 1 {
t.Fatalf("limit=%d: expected 1 result but got %v", limit, resp.Results)
}
row, ok := resp.Results[0].(*pilosa.Row)
if !ok {
t.Fatalf("limit=%d: expected a row result but got %T", limit, resp.Results[0])
}
got := row.Columns()
if !reflect.DeepEqual(expect, got) {
t.Errorf("limit=%d: expected %v but got %v", limit, expect, got)
}
}
})
// Test with only an offset specified.
t.Run("Offset", func(t *testing.T) {
for offset := 0; offset < 5; offset++ {
expect := []uint64{}
if offset <= len(columns) {
expect = columns[offset:]
}
resp := c.Query(t, "i", fmt.Sprintf("Limit(All(), offset=%d)", offset))
if len(resp.Results) != 1 {
t.Fatalf("offset=%d: expected 1 result but got %v", offset, resp.Results)
}
row, ok := resp.Results[0].(*pilosa.Row)
if !ok {
t.Fatalf("offset=%d: expected a row result but got %T", offset, resp.Results[0])
}
got := row.Columns()
if !reflect.DeepEqual(expect, got) {
t.Errorf("offset=%d: expected %v but got %v", offset, expect, got)
}
}
})
// Test with a limit and offset specified.
t.Run("LimitOffset", func(t *testing.T) {
for limit := 0; limit < 5; limit++ {
for offset := 0; offset < 5; offset++ {
expect := []uint64{}
if offset <= len(columns) {
expect = columns[offset:]
}
if limit < len(expect) {
expect = expect[:limit]
}
resp := c.Query(t, "i", fmt.Sprintf("Limit(All(), limit=%d, offset=%d)", limit, offset))
if len(resp.Results) != 1 {
t.Fatalf("limit=%d,offset=%d: expected 1 result but got %v", limit, offset, resp.Results)
}
row, ok := resp.Results[0].(*pilosa.Row)
if !ok {
t.Fatalf("limit=%d,offset=%d: expected a row result but got %T", limit, offset, resp.Results[0])
}
got := row.Columns()
if !reflect.DeepEqual(expect, got) {
t.Errorf("limit=%d,offset=%d: expected %v but got %v", limit, offset, expect, got)
}
}
}
})
}
// Ensure an all query can be executed.
func TestExecutor_Execute_All(t *testing.T) {
t.Run("ColumnID", func(t *testing.T) {

View file

@ -394,7 +394,14 @@ var callInfoByFunc = map[string]callInfo{
"Union": {allowUnknown: false},
"UnionRows": {allowUnknown: false},
"Extract": {allowUnknown: false},
"Xor": {allowUnknown: false},
"Limit": {
allowUnknown: false,
prototypes: map[string]interface{}{
"limit": int64(0),
"offset": int64(0),
},
},
"Xor": {allowUnknown: false},
"ConstRow": {
allowUnknown: false,