planner: support push window function down to tiflash

executor/partition_table.go

@@ -52,6 +52,10 @@ func updateExecutorTableID(ctx context.Context, exec *tipb.Executor, recursive b
 		child = exec.Join.Children[1-exec.Join.InnerIdx]
 	case tipb.ExecType_TypeProjection:
 		child = exec.Projection.Child
+	case tipb.ExecType_TypeWindow:
+		child = exec.Window.Child
+	case tipb.ExecType_TypeSort:
+		child = exec.Sort.Child
 	default:
 		return errors.Trace(fmt.Errorf("unknown new tipb protocol %d", exec.Tp))
 	}

expression/aggregation/agg_to_pb.go

@@ -28,11 +28,9 @@ import (
 	"github.com/pingcap/tipb/go-tipb"
 )
 
-// AggFuncToPBExpr converts aggregate function to pb.
-func AggFuncToPBExpr(sctx sessionctx.Context, client kv.Client, aggFunc *AggFuncDesc) *tipb.Expr {
-	pc := expression.NewPBConverter(client, sctx.GetSessionVars().StmtCtx)
-	var tp tipb.ExprType
-	switch aggFunc.Name {
+// GetTiPBExpr return the TiPB ExprType of desc.
+func (desc *baseFuncDesc) GetTiPBExpr(tryWindowDesc bool) (tp tipb.ExprType) {
+	switch desc.Name {
 	case ast.AggFuncCount:
 		tp = tipb.ExprType_Count
 	case ast.AggFuncApproxCountDistinct:
@@ -68,6 +66,42 @@ func AggFuncToPBExpr(sctx sessionctx.Context, client kv.Client, aggFunc *AggFunc
 	case ast.AggFuncStddevSamp:
 		tp = tipb.ExprType_StddevSamp
 	}
+
+	if tp != tipb.ExprType_Null || !tryWindowDesc {
+		return
+	}
+
+	switch desc.Name {
+	case ast.WindowFuncRowNumber:
+		tp = tipb.ExprType_RowNumber
+	case ast.WindowFuncRank:
+		tp = tipb.ExprType_Rank
+	case ast.WindowFuncDenseRank:
+		tp = tipb.ExprType_DenseRank
+	case ast.WindowFuncCumeDist:
+		tp = tipb.ExprType_CumeDist
+	case ast.WindowFuncPercentRank:
+		tp = tipb.ExprType_PercentRank
+	case ast.WindowFuncNtile:
+		tp = tipb.ExprType_Ntile
+	case ast.WindowFuncLead:
+		tp = tipb.ExprType_Lead
+	case ast.WindowFuncLag:
+		tp = tipb.ExprType_Lag
+	case ast.WindowFuncFirstValue:
+		tp = tipb.ExprType_FirstValue
+	case ast.WindowFuncLastValue:
+		tp = tipb.ExprType_LastValue
+	case ast.WindowFuncNthValue:
+		tp = tipb.ExprType_NthValue
+	}
+	return tp
+}
+
+// AggFuncToPBExpr converts aggregate function to pb.
+func AggFuncToPBExpr(sctx sessionctx.Context, client kv.Client, aggFunc *AggFuncDesc) *tipb.Expr {
+	pc := expression.NewPBConverter(client, sctx.GetSessionVars().StmtCtx)
+	tp := aggFunc.GetTiPBExpr(false)
 	if !client.IsRequestTypeSupported(kv.ReqTypeSelect, int64(tp)) {
 		return nil
 	}

expression/aggregation/window_func.go

@@ -18,8 +18,10 @@ import (
 	"strings"
 
 	"github.com/pingcap/tidb/expression"
+	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/sessionctx"
+	"github.com/pingcap/tipb/go-tipb"
 )
 
 // WindowFuncDesc describes a window function signature, only used in planner.
@@ -92,3 +94,41 @@ func NeedFrame(name string) bool {
 	_, ok := noFrameWindowFuncs[strings.ToLower(name)]
 	return !ok
 }
+
+// Clone makes a copy of SortItem.
+func (s *WindowFuncDesc) Clone() *WindowFuncDesc {
+	return &WindowFuncDesc{*s.baseFuncDesc.clone()}
+}
+
+// WindowFuncToPBExpr converts aggregate function to pb.
+func WindowFuncToPBExpr(sctx sessionctx.Context, client kv.Client, desc *WindowFuncDesc) *tipb.Expr {
+	pc := expression.NewPBConverter(client, sctx.GetSessionVars().StmtCtx)
+	tp := desc.GetTiPBExpr(true)
+	if !client.IsRequestTypeSupported(kv.ReqTypeSelect, int64(tp)) {
+		return nil
+	}
+
+	children := make([]*tipb.Expr, 0, len(desc.Args))
+	for _, arg := range desc.Args {
+		pbArg := pc.ExprToPB(arg)
+		if pbArg == nil {
+			return nil
+		}
+		children = append(children, pbArg)
+	}
+	return &tipb.Expr{Tp: tp, Children: children, FieldType: expression.ToPBFieldType(desc.RetTp)}
+}
+
+// CanPushDownToTiFlash control whether a window function desc can be push down to tiflash.
+func (s *WindowFuncDesc) CanPushDownToTiFlash() bool {
+	// window functions
+	switch s.Name {
+	case ast.WindowFuncRowNumber, ast.WindowFuncRank, ast.WindowFuncDenseRank, ast.WindowFuncLead, ast.WindowFuncLag:
+		return true
+		// TODO: support aggregate functions
+		//case ast.AggFuncSum, ast.AggFuncCount, ast.AggFuncAvg, ast.AggFuncMax, ast.AggFuncMin:
+		//	return true
+	}
+
+	return false
+}

go.mod

@@ -52,7 +52,7 @@ require (
 	github.com/pingcap/log v1.1.0
 	github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4
 	github.com/pingcap/tidb/parser v0.0.0-20211011031125-9b13dc409c5e
-	github.com/pingcap/tipb v0.0.0-20220215045658-d12dec7a7609
+	github.com/pingcap/tipb v0.0.0-20220314125451-bfb5c2c55188
 	github.com/prometheus/client_golang v1.11.0
 	github.com/prometheus/client_model v0.2.0
 	github.com/prometheus/common v0.32.1

go.sum

@@ -643,8 +643,8 @@ github.com/pingcap/log v1.1.0 h1:ELiPxACz7vdo1qAvvaWJg1NrYFoY6gqAh/+Uo6aXdD8=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4 h1:HYbcxtnkN3s5tqrZ/z3eJS4j3Db8wMphEm1q10lY/TM=
 github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4/go.mod h1:sDCsM39cGiv2vwunZkaFA917vVkqDTGSPbbV7z4Oops=
-github.com/pingcap/tipb v0.0.0-20220215045658-d12dec7a7609 h1:BiCS1ZRnW0szOvTAa3gCqWIhyo+hv83SVaBgrUghXIU=
-github.com/pingcap/tipb v0.0.0-20220215045658-d12dec7a7609/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
+github.com/pingcap/tipb v0.0.0-20220314125451-bfb5c2c55188 h1:+46isFI9fR9R+nJVDMI55tCC/TCwp+bvVA4HLGEv1rY=
+github.com/pingcap/tipb v0.0.0-20220314125451-bfb5c2c55188/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4 h1:49lOXmGaUpV9Fz3gd7TFZY106KVlPVa5jcYD1gaQf98=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4/go.mod h1:4OwLy04Bl9Ef3GJJCoec+30X3LQs/0/m4HFRt/2LUSA=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

kv/checker.go

@@ -14,7 +14,9 @@
 
 package kv
 
-import "github.com/pingcap/tipb/go-tipb"
+import (
+	"github.com/pingcap/tipb/go-tipb"
+)
 
 // RequestTypeSupportedChecker is used to check expression can be pushed down.
 type RequestTypeSupportedChecker struct{}
@@ -37,6 +39,10 @@ func (d RequestTypeSupportedChecker) IsRequestTypeSupported(reqType, subType int
 	return false
 }
 
+// TODO: deprecate this function, because:
+// 1. we have more than one storage engine available for push-down.
+// 2. we'd better do an accurate push-down check at the planner stage, instead of here.
+// currently, we do aggregation push-down check in `CheckAggCanPushCop`.
 func (d RequestTypeSupportedChecker) supportExpr(exprType tipb.ExprType) bool {
 	switch exprType {
 	case tipb.ExprType_Null, tipb.ExprType_Int64, tipb.ExprType_Uint64, tipb.ExprType_String, tipb.ExprType_Bytes,
@@ -48,6 +54,10 @@ func (d RequestTypeSupportedChecker) supportExpr(exprType tipb.ExprType) bool {
 	case tipb.ExprType_Count, tipb.ExprType_First, tipb.ExprType_Max, tipb.ExprType_Min, tipb.ExprType_Sum, tipb.ExprType_Avg,
 		tipb.ExprType_Agg_BitXor, tipb.ExprType_Agg_BitAnd, tipb.ExprType_Agg_BitOr, tipb.ExprType_ApproxCountDistinct, tipb.ExprType_GroupConcat:
 		return true
+	// window functions.
+	case tipb.ExprType_RowNumber, tipb.ExprType_Rank, tipb.ExprType_DenseRank, tipb.ExprType_CumeDist, tipb.ExprType_PercentRank,
+		tipb.ExprType_Ntile, tipb.ExprType_Lead, tipb.ExprType_Lag, tipb.ExprType_FirstValue, tipb.ExprType_LastValue, tipb.ExprType_NthValue:
+		return true
 	case ReqSubTypeDesc:
 		return true
 	case ReqSubTypeSignature:

planner/cascades/enforcer_rules.go

@@ -41,7 +41,7 @@ func GetEnforcerRules(g *memo.Group, prop *property.PhysicalProperty) (enforcers
 	if g.EngineType != memo.EngineTiDB {
 		return
 	}
-	if !prop.IsEmpty() {
+	if !prop.IsSortItemEmpty() {
 		enforcers = append(enforcers, orderEnforcer)
 	}
 	return

planner/cascades/implementation_rules.go

@@ -96,7 +96,7 @@ type ImplTableDual struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplTableDual) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -114,7 +114,7 @@ type ImplMemTableScan struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplMemTableScan) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -189,15 +189,15 @@ type ImplTableScan struct {
 // Match implements ImplementationRule Match interface.
 func (r *ImplTableScan) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
 	ts := expr.ExprNode.(*plannercore.LogicalTableScan)
-	return prop.IsEmpty() || (len(prop.SortItems) == 1 && ts.HandleCols != nil && prop.SortItems[0].Col.Equal(nil, ts.HandleCols.GetCol(0)))
+	return prop.IsSortItemEmpty() || (len(prop.SortItems) == 1 && ts.HandleCols != nil && prop.SortItems[0].Col.Equal(nil, ts.HandleCols.GetCol(0)))
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
 func (r *ImplTableScan) OnImplement(expr *memo.GroupExpr, reqProp *property.PhysicalProperty) ([]memo.Implementation, error) {
 	logicProp := expr.Group.Prop
 	logicalScan := expr.ExprNode.(*plannercore.LogicalTableScan)
 	ts := logicalScan.GetPhysicalScan(logicProp.Schema, logicProp.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt))
-	if !reqProp.IsEmpty() {
+	if !reqProp.IsSortItemEmpty() {
 		ts.KeepOrder = true
 		ts.Desc = reqProp.SortItems[0].Desc
 	}
@@ -219,7 +219,7 @@ func (r *ImplIndexScan) Match(expr *memo.GroupExpr, prop *property.PhysicalPrope
 func (r *ImplIndexScan) OnImplement(expr *memo.GroupExpr, reqProp *property.PhysicalProperty) ([]memo.Implementation, error) {
 	logicalScan := expr.ExprNode.(*plannercore.LogicalIndexScan)
 	is := logicalScan.GetPhysicalIndexScan(expr.Group.Prop.Schema, expr.Group.Prop.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt))
-	if !reqProp.IsEmpty() {
+	if !reqProp.IsSortItemEmpty() {
 		is.KeepOrder = true
 		if reqProp.SortItems[0].Desc {
 			is.Desc = true
@@ -235,7 +235,7 @@ type ImplShow struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplShow) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -319,7 +319,7 @@ type ImplHashAgg struct {
 // Match implements ImplementationRule Match interface.
 func (r *ImplHashAgg) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
 	// TODO: deal with the hints when we have implemented StreamAgg.
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -348,7 +348,7 @@ type ImplLimit struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplLimit) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -371,7 +371,7 @@ type ImplTopN struct {
 func (r *ImplTopN) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
 	topN := expr.ExprNode.(*plannercore.LogicalTopN)
 	if expr.Group.EngineType != memo.EngineTiDB {
-		return prop.IsEmpty()
+		return prop.IsSortItemEmpty()
 	}
 	return plannercore.MatchItems(prop, topN.ByItems)
 }
@@ -446,7 +446,7 @@ type ImplHashJoinBuildLeft struct {
 func (r *ImplHashJoinBuildLeft) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
 	switch expr.ExprNode.(*plannercore.LogicalJoin).JoinType {
 	case plannercore.InnerJoin, plannercore.LeftOuterJoin, plannercore.RightOuterJoin:
-		return prop.IsEmpty()
+		return prop.IsSortItemEmpty()
 	default:
 		return false
 	}
@@ -473,7 +473,7 @@ type ImplHashJoinBuildRight struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplHashJoinBuildRight) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -520,7 +520,7 @@ type ImplUnionAll struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplUnionAll) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface.
@@ -572,7 +572,7 @@ type ImplMaxOneRow struct {
 
 // Match implements ImplementationRule Match interface.
 func (r *ImplMaxOneRow) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
-	return prop.IsEmpty()
+	return prop.IsSortItemEmpty()
 }
 
 // OnImplement implements ImplementationRule OnImplement interface