func (h *rpcHandler) handleCopRequest(req *coprocessor.Request) (*coprocessor.Response, error) { resp := &coprocessor.Response{} if err := h.checkContext(req.GetContext()); err != nil { resp.RegionError = err return resp, nil } if len(req.Ranges) == 0 { return resp, nil } if req.GetTp() == kv.ReqTypeSelect || req.GetTp() == kv.ReqTypeIndex { sel := new(tipb.SelectRequest) err := proto.Unmarshal(req.Data, sel) if err != nil { return nil, errors.Trace(err) } ctx := &selectContext{ sel: sel, keyRanges: req.Ranges, sc: xeval.FlagsToStatementContext(sel.Flags), } ctx.eval = xeval.NewEvaluator(ctx.sc) if sel.Where != nil { ctx.whereColumns = make(map[int64]*tipb.ColumnInfo) collectColumnsInExpr(sel.Where, ctx, ctx.whereColumns) } ctx.aggregate = len(sel.Aggregates) > 0 || len(sel.GetGroupBy()) > 0 if ctx.aggregate { // compose aggregateFuncExpr ctx.aggregates = make([]*aggregateFuncExpr, 0, len(sel.Aggregates)) ctx.aggColumns = make(map[int64]*tipb.ColumnInfo) for _, agg := range sel.Aggregates { aggExpr := &aggregateFuncExpr{expr: agg} ctx.aggregates = append(ctx.aggregates, aggExpr) collectColumnsInExpr(agg, ctx, ctx.aggColumns) } ctx.groups = make(map[string]bool) ctx.groupKeys = make([][]byte, 0) for _, item := range ctx.sel.GetGroupBy() { collectColumnsInExpr(item.Expr, ctx, ctx.aggColumns) } for k := range ctx.whereColumns { // It is will be handled in where. delete(ctx.aggColumns, k) } } var chunks []tipb.Chunk if req.GetTp() == kv.ReqTypeSelect { chunks, err = h.getChunksFromSelectReq(ctx) } else { // The PKHandle column info has been collected in ctx, so we can remove it in IndexInfo. length := len(sel.IndexInfo.Columns) if sel.IndexInfo.Columns[length-1].GetPkHandle() { sel.IndexInfo.Columns = sel.IndexInfo.Columns[:length-1] } chunks, err = h.getChunksFromIndexReq(ctx) } selResp := new(tipb.SelectResponse) selResp.Error = toPBError(err) selResp.Chunks = chunks if err != nil { if locked, ok := errors.Cause(err).(*ErrLocked); ok { resp.Locked = &kvrpcpb.LockInfo{ Key: locked.Key, PrimaryLock: locked.Primary, LockVersion: locked.StartTS, LockTtl: locked.TTL, } } else { resp.OtherError = err.Error() } } data, err := proto.Marshal(selResp) if err != nil { return nil, errors.Trace(err) } resp.Data = data } return resp, nil }
func (rs *localRegion) Handle(req *regionRequest) (*regionResponse, error) { resp := ®ionResponse{ req: req, } if req.Tp == kv.ReqTypeSelect || req.Tp == kv.ReqTypeIndex { sel := new(tipb.SelectRequest) err := proto.Unmarshal(req.data, sel) if err != nil { return nil, errors.Trace(err) } txn := newTxn(rs.store, kv.Version{Ver: uint64(sel.StartTs)}) ctx := &selectContext{ sel: sel, txn: txn, keyRanges: req.ranges, sc: xeval.FlagsToStatementContext(sel.Flags), } ctx.eval = xeval.NewEvaluator(ctx.sc) if sel.Where != nil { ctx.whereColumns = make(map[int64]*tipb.ColumnInfo) collectColumnsInExpr(sel.Where, ctx, ctx.whereColumns) } if len(sel.OrderBy) > 0 { if sel.OrderBy[0].Expr == nil { ctx.descScan = sel.OrderBy[0].Desc } else { if sel.Limit == nil { return nil, errors.New("we don't support pushing down Sort without Limit") } ctx.topn = true ctx.topnHeap = &topnHeap{ totalCount: int(*sel.Limit), topnSorter: topnSorter{ orderByItems: sel.OrderBy, ctx: ctx, }, } ctx.topnColumns = make(map[int64]*tipb.ColumnInfo) for _, item := range sel.OrderBy { collectColumnsInExpr(item.Expr, ctx, ctx.topnColumns) } for k := range ctx.whereColumns { // It will be handled in where. delete(ctx.topnColumns, k) } } } ctx.aggregate = len(sel.Aggregates) > 0 || len(sel.GetGroupBy()) > 0 if ctx.aggregate { // compose aggregateFuncExpr ctx.aggregates = make([]*aggregateFuncExpr, 0, len(sel.Aggregates)) ctx.aggColumns = make(map[int64]*tipb.ColumnInfo) for _, agg := range sel.Aggregates { aggExpr := &aggregateFuncExpr{expr: agg} ctx.aggregates = append(ctx.aggregates, aggExpr) collectColumnsInExpr(agg, ctx, ctx.aggColumns) } ctx.groups = make(map[string]bool) ctx.groupKeys = make([][]byte, 0) for _, item := range ctx.sel.GetGroupBy() { collectColumnsInExpr(item.Expr, ctx, ctx.aggColumns) } for k := range ctx.whereColumns { // It will be handled in where. delete(ctx.aggColumns, k) } } if req.Tp == kv.ReqTypeSelect { err = rs.getRowsFromSelectReq(ctx) } else { // The PKHandle column info has been collected in ctx, so we can remove it in IndexInfo. length := len(sel.IndexInfo.Columns) if sel.IndexInfo.Columns[length-1].GetPkHandle() { sel.IndexInfo.Columns = sel.IndexInfo.Columns[:length-1] } err = rs.getRowsFromIndexReq(ctx) } if ctx.topn { rs.setTopNDataForCtx(ctx) } selResp := new(tipb.SelectResponse) selResp.Error = toPBError(err) selResp.Chunks = ctx.chunks resp.err = err data, err := proto.Marshal(selResp) if err != nil { return nil, errors.Trace(err) } resp.data = data } if bytes.Compare(rs.startKey, req.startKey) < 0 || bytes.Compare(rs.endKey, req.endKey) > 0 { resp.newStartKey = rs.startKey resp.newEndKey = rs.endKey } return resp, nil }