diff --git a/executor.go b/executor.go index 0a5b207ae..774f137fe 100644 --- a/executor.go +++ b/executor.go @@ -565,130 +565,6 @@ func (e *executor) execute(ctx context.Context, qcx *Qcx, index string, q *pql.Q // preprocessQuery expands any calls that need preprocessing. func (e *executor) preprocessQuery(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*pql.Call, error) { switch c.Name { - case "UnionRows": - // Turn UnionRows(Rows(...)) into Union(Row(...), ...). - var rows []*pql.Call - for _, child := range c.Children { - // Check that we can use the call. - switch child.Name { - case "Rows": - case "TopN": - default: - return nil, errors.Errorf("cannot use %v as a rows query", child) - } - - // Execute the call. - rowsResult, err := e.executeCall(ctx, qcx, index, child, shards, opt) - if err != nil { - return nil, err - } - - // Turn the results into rows calls. - var resultRows []*pql.Call - switch rowsResult := rowsResult.(type) { - case *PairsField: - // Translate pairs into rows calls. - for _, p := range rowsResult.Pairs { - var val interface{} - switch { - case p.Key != "": - val = p.Key - default: - val = p.ID - } - resultRows = append(resultRows, &pql.Call{ - Name: "Row", - Args: map[string]interface{}{ - rowsResult.Field: val, - }, - }) - } - 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, - }, - }) - } - default: - return nil, errors.Errorf("unexpected Rows type %T", rowsResult) - } - - // Propogate any special properties of the call. - switch child.Name { - case "Rows": - // Propogate "from" time, if set. - if v, ok := child.Args["from"]; ok { - for _, rowCall := range resultRows { - rowCall.Args["from"] = v - } - } - - // Propogate "to" time, if set. - if v, ok := child.Args["to"]; ok { - for _, rowCall := range resultRows { - rowCall.Args["to"] = v - } - } - } - - rows = append(rows, resultRows...) - } - - // Generate a Union call over the rows. - return &pql.Call{ - Name: "Union", - Children: rows, - }, nil - - case "ConstRow": - // Fetch user-provided columns list. - cols, _ := c.Args["columns"].([]interface{}) - var ids []uint64 - var keys []string - for _, c := range cols { - switch c := c.(type) { - case uint64: - ids = append(ids, c) - case int64: - ids = append(ids, uint64(c)) - case string: - keys = append(keys, c) - default: - return nil, errors.Errorf("invalid column identifier %v of type %T", c, c) - } - } - - // Translate keys to IDs. - if len(keys) > 0 { - keyIDs, err := e.Cluster.translateIndexKeys(ctx, index, keys, false) - if err != nil { - return nil, errors.Wrap(err, "translating column IDs in ConstRow") - } - ids = append(ids, keyIDs...) - } - - // Split IDs by shard. - shardSet := make(map[uint64][]uint64) - for _, id := range ids { - shardSet[id/ShardWidth] = append(shardSet[id/ShardWidth], id) - } - - // Convert ID sets to per-shard Row objects. - precomputed := make(map[uint64]interface{}) - for _, s := range shards { - precomputed[s] = NewRow(shardSet[s]...) - } - - // Generate a precomputed call with the data. - return &pql.Call{ - Name: "Precomputed", - Precomputed: precomputed, - }, nil - case "All": _, hasLimit, err := c.UintArg("limit") if err != nil { @@ -709,21 +585,6 @@ func (e *executor) preprocessQuery(ctx context.Context, qcx *Qcx, index string, }, } 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, qcx, index, c.Children[0], shards, opt) - if err != nil { - return nil, err - } - c.Children[0] = res - err = e.executeLimitCall(ctx, qcx, index, c, shards, opt) - if err != nil { - return nil, err - } return c, nil default: @@ -855,6 +716,12 @@ func (e *executor) executeCall(ctx context.Context, qcx *Qcx, index string, c *p return e.executeFieldValueCall(ctx, qcx, index, c, shards, opt) case "Precomputed": return e.executePrecomputedCall(ctx, qcx, index, c, shards, opt) + case "UnionRows": + return e.executeUnionRows(ctx, qcx, index, c, shards, opt) + case "ConstRow": + return e.executeConstRow(ctx, index, c) + case "Limit": + return e.executeLimitCall(ctx, qcx, index, c, shards, opt) default: // e.g. "Row", "Union", "Intersect" or anything that returns a bitmap. statFn() return e.executeBitmapCall(ctx, qcx, index, c, shards, opt) @@ -1055,100 +922,80 @@ func (e *executor) executeFieldValueCallShard(ctx context.Context, qcx *Qcx, fie return other, nil } -// executeLimitCall executes a Limit() call, **rewriting it to a precomputed call**. -func (e *executor) executeLimitCall(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) error { +// executeLimitCall executes a Limit() call. +func (e *executor) executeLimitCall(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*Row, error) { bitmapCall := c.Children[0] limit, hasLimit, err := c.UintArg("limit") if err != nil { - return errors.Wrap(err, "getting limit") + return nil, errors.Wrap(err, "getting limit") } offset, _, err := c.UintArg("offset") if err != nil { - return errors.Wrap(err, "getting offset") + return nil, errors.Wrap(err, "getting offset") } if !hasLimit { limit = math.MaxUint64 } - // skip tracks the number of records left to be skipped - // in support of getting to the offset. - var skip uint64 = offset - - // got tracks the number of records gotten to that point. - var got uint64 - - c.Precomputed = make(map[uint64]interface{}) - - for _, shard := range shards { - // Execute calls in bulk on each remote node and merge. - mapFn := func(ctx context.Context, shard uint64) (_ interface{}, err error) { - return e.executeBitmapCallShard(ctx, qcx, 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. - if segCnt <= skip { - skip -= segCnt - continue - } - - // This segment doesn't have enough to finish fulfilling the limit - // (or it has exactly enough). - if segCnt-skip <= limit-got { - if skip == 0 { - c.Precomputed[shard] = row - } else { - cols := row.Columns() - partialRow := NewRow() - for _, bit := range cols[skip:] { - partialRow.SetBit(bit) - } - c.Precomputed[shard] = partialRow - } - got += segCnt - skip - // In the case where this segment exactly fulfills the limit, break. - if got == limit { - break - } - skip = 0 - continue - } - - // This segment has more records than the remaining limit requires. - cols := row.Columns() - partialRow := NewRow() - for _, bit := range cols[skip : skip+limit-got] { - partialRow.SetBit(bit) - } - c.Precomputed[shard] = partialRow - break + // Execute bitmap call, storing the full result on this node. + res, err := e.executeCall(ctx, qcx, index, bitmapCall, shards, opt) + if err != nil { + return nil, errors.Wrap(err, "limit map reduce") + } + if res == nil { + res = NewRow() } - c.Name = "Precomputed" - return nil + result, ok := res.(*Row) + if !ok { + return nil, errors.Errorf("expected Row but got %T", result) + } + + if offset != 0 { + i := 0 + var leadingBits []uint64 + for i < len(result.segments) && offset > 0 { + seg := result.segments[i] + count := seg.Count() + if count > offset { + data := seg.Columns() + data = data[offset:] + leadingBits = data + i++ + break + } + + offset -= count + i++ + } + row := NewRow(leadingBits...) + row.Merge(&Row{segments: result.segments[i:]}) + result = row + } + if limit < result.Count() { + i := 0 + var trailingBits []uint64 + for i < len(result.segments) && limit > 0 { + seg := result.segments[i] + count := seg.Count() + if count > limit { + data := seg.Columns() + data = data[:limit] + trailingBits = data + break + } + + limit -= count + i++ + } + row := NewRow(trailingBits...) + row.Merge(&Row{segments: result.segments[:i]}) + result = row + } + + return result, nil } // executeIncludesColumnCallShard @@ -3694,6 +3541,119 @@ func (e *executor) executeNotShard(ctx context.Context, qcx *Qcx, index string, return existenceRow.Difference(row), nil } +func (e *executor) executeConstRow(ctx context.Context, index string, c *pql.Call) (res *Row, err error) { + // Fetch user-provided columns list. + cols, _ := c.Args["columns"].([]interface{}) + var ids []uint64 + var keys []string + for _, c := range cols { + switch c := c.(type) { + case uint64: + ids = append(ids, c) + case int64: + ids = append(ids, uint64(c)) + case string: + keys = append(keys, c) + default: + return nil, errors.Errorf("invalid column identifier %v of type %T", c, c) + } + } + + // Translate keys to IDs. + if len(keys) > 0 { + keyIDs, err := e.Cluster.translateIndexKeys(ctx, index, keys, false) + if err != nil { + return nil, errors.Wrap(err, "translating column IDs in ConstRow") + } + ids = append(ids, keyIDs...) + } + + return NewRow(ids...), nil +} + +func (e *executor) executeUnionRows(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*Row, error) { + // Turn UnionRows(Rows(...)) into Union(Row(...), ...). + var rows []*pql.Call + for _, child := range c.Children { + // Check that we can use the call. + switch child.Name { + case "Rows": + case "TopN": + default: + return nil, errors.Errorf("cannot use %v as a rows query", child) + } + + // Execute the call. + rowsResult, err := e.executeCall(ctx, qcx, index, child, shards, opt) + if err != nil { + return nil, err + } + + // Turn the results into rows calls. + var resultRows []*pql.Call + switch rowsResult := rowsResult.(type) { + case *PairsField: + // Translate pairs into rows calls. + for _, p := range rowsResult.Pairs { + var val interface{} + switch { + case p.Key != "": + val = p.Key + default: + val = p.ID + } + resultRows = append(resultRows, &pql.Call{ + Name: "Row", + Args: map[string]interface{}{ + rowsResult.Field: val, + }, + }) + } + 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, + }, + }) + } + default: + return nil, errors.Errorf("unexpected Rows type %T", rowsResult) + } + + // Propogate any special properties of the call. + switch child.Name { + case "Rows": + // Propogate "from" time, if set. + if v, ok := child.Args["from"]; ok { + for _, rowCall := range resultRows { + rowCall.Args["from"] = v + } + } + + // Propogate "to" time, if set. + if v, ok := child.Args["to"]; ok { + for _, rowCall := range resultRows { + rowCall.Args["to"] = v + } + } + } + + rows = append(rows, resultRows...) + } + + // Generate a Union call over the rows. + c = &pql.Call{ + Name: "Union", + Children: rows, + } + + // Execute the generated Union() call. + return e.executeBitmapCall(ctx, qcx, index, c, shards, opt) +} + // executeAllCallShard executes an All() call for a local shard. func (e *executor) executeAllCallShard(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shard uint64) (res *Row, err error) { diff --git a/executor_test.go b/executor_test.go index da04e4117..c44df08f1 100644 --- a/executor_test.go +++ b/executor_test.go @@ -3659,7 +3659,7 @@ 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) + c := test.MustRunCluster(t, 2) defer c.Close() c.CreateField(t, "i", pilosa.IndexOptions{TrackExistence: true}, "f") @@ -3743,6 +3743,58 @@ func TestExecutor_Execute_Limit(t *testing.T) { } } }) + + t.Run("Nested", 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(Limit(All(), offset=%d), limit=%d)", offset, limit)) + 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) + } + } + } + }) + + t.Run("Extract", func(t *testing.T) { + resp := c.Query(t, "i", "Extract(Limit(All(), limit=1))") + if len(resp.Results) != 1 { + t.Fatalf("expected 1 result but got %d", len(resp.Results)) + } + got, ok := resp.Results[0].(pilosa.ExtractedTable) + if !ok { + t.Fatalf("expected a table result but got %T", resp.Results[0]) + } + expect := pilosa.ExtractedTable{ + Fields: []pilosa.ExtractedTableField{}, + Columns: []pilosa.ExtractedTableColumn{ + { + Column: pilosa.KeyOrID{ + ID: 0, + }, + Rows: []interface{}{}, + }, + }, + } + if !reflect.DeepEqual(expect, got) { + t.Errorf("expected %v but got %v", expect, got) + } + }) } // Ensure an all query can be executed. diff --git a/pql/ast.go b/pql/ast.go index db142fcd0..b2e03fece 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -392,7 +392,7 @@ var callInfoByFunc = map[string]callInfo{ }, }, "Union": {allowUnknown: false}, - "UnionRows": {allowUnknown: false}, + "UnionRows": {allowUnknown: false, callType: PrecallGlobal}, "Extract": {allowUnknown: false}, "Limit": { allowUnknown: false, @@ -400,6 +400,7 @@ var callInfoByFunc = map[string]callInfo{ "limit": int64(0), "offset": int64(0), }, + callType: PrecallGlobal, }, "Xor": {allowUnknown: false}, @@ -408,6 +409,7 @@ var callInfoByFunc = map[string]callInfo{ prototypes: map[string]interface{}{ "columns": []interface{}{}, }, + callType: PrecallGlobal, }, // things that take _field