planner/cascades: Add ImplementationRule for Window (#14085)

planner/cascades/implementation_rules.go

@@ -80,6 +80,9 @@ var defaultImplementationMap = map[memo.Operand][]ImplementationRule{
 	memo.OperandMaxOneRow: {
 		&ImplMaxOneRow{},
 	},
+	memo.OperandWindow: {
+		&ImplWindow{},
+	},
 }
 
 // ImplTableDual implements LogicalTableDual as PhysicalTableDual.
@@ -519,3 +522,38 @@ func (r *ImplMaxOneRow) OnImplement(expr *memo.GroupExpr, reqProp *property.Phys
 		&property.PhysicalProperty{ExpectedCnt: 2})
 	return impl.NewMaxOneRowImpl(physicalMaxOneRow), nil
 }
+
+// ImplWindow implements LogicalWindow to PhysicalWindow.
+type ImplWindow struct {
+}
+
+// Match implements ImplementationRule Match interface.
+func (w *ImplWindow) Match(expr *memo.GroupExpr, prop *property.PhysicalProperty) (matched bool) {
+	lw := expr.ExprNode.(*plannercore.LogicalWindow)
+	var byItems []property.Item
+	byItems = append(byItems, lw.PartitionBy...)
+	byItems = append(byItems, lw.OrderBy...)
+	childProperty := &property.PhysicalProperty{ExpectedCnt: math.MaxFloat64, Items: byItems}
+	return prop.IsPrefix(childProperty)
+}
+
+// OnImplement implements ImplementationRule OnImplement interface.
+func (w *ImplWindow) OnImplement(expr *memo.GroupExpr, reqProp *property.PhysicalProperty) (memo.Implementation, error) {
+	lw := expr.ExprNode.(*plannercore.LogicalWindow)
+	var byItems []property.Item
+	byItems = append(byItems, lw.PartitionBy...)
+	byItems = append(byItems, lw.OrderBy...)
+	physicalWindow := plannercore.PhysicalWindow{
+		WindowFuncDescs: lw.WindowFuncDescs,
+		PartitionBy:     lw.PartitionBy,
+		OrderBy:         lw.OrderBy,
+		Frame:           lw.Frame,
+	}.Init(
+		lw.SCtx(),
+		expr.Group.Prop.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt),
+		lw.SelectBlockOffset(),
+		&property.PhysicalProperty{ExpectedCnt: math.MaxFloat64, Items: byItems},
+	)
+	physicalWindow.SetSchema(expr.Group.Prop.Schema)
+	return impl.NewWindowImpl(physicalWindow), nil
+}

planner/cascades/testdata/integration_suite_in.json

@@ -38,7 +38,8 @@
       "select a from t limit 1 offset 2",
       "select b from t order by b limit 3",
       "select a from t order by a limit 1 offset 2",
-      "select * from ((select a as aa from t t1) union all (select b as aa from t t2)) as t3 order by aa"
+      "select * from ((select a as aa from t t1) union all (select b as aa from t t2)) as t3 order by aa",
+      "select a, b, lag(a,1) over (order by b) from t order by b"
     ]
   },
   {

planner/cascades/testdata/integration_suite_out.json

@@ -312,6 +312,21 @@
           "33",
           "44"
         ]
+      },
+      {
+        "SQL": "select a, b, lag(a,1) over (order by b) from t order by b",
+        "Plan": [
+          "Window_10 10000.00 root lag(test.t.a, 1)->Column#4 over(order by test.t.b asc)",
+          "└─Sort_15 10000.00 root test.t.b:asc",
+          "  └─TableReader_13 10000.00 root data:TableScan_14",
+          "    └─TableScan_14 10000.00 cop[tikv] table:t, range:[-inf,+inf], keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 11 <nil>",
+          "2 22 1",
+          "3 33 2",
+          "4 44 3"
+        ]
       }
     ]
   },

planner/implementation/simple_plans.go

@@ -203,3 +203,19 @@ func (impl *MaxOneRowImpl) CalcCost(outCount float64, children ...memo.Implement
 func NewMaxOneRowImpl(maxOneRow *plannercore.PhysicalMaxOneRow) *MaxOneRowImpl {
 	return &MaxOneRowImpl{baseImpl{plan: maxOneRow}}
 }
+
+// WindowImpl is the implementation of PhysicalWindow.
+type WindowImpl struct {
+	baseImpl
+}
+
+// NewWindowImpl creates a new WindowImpl.
+func NewWindowImpl(window *plannercore.PhysicalWindow) *WindowImpl {
+	return &WindowImpl{baseImpl{plan: window}}
+}
+
+// CalcCost implements Implementation CalcCost interface.
+func (impl *WindowImpl) CalcCost(outCount float64, children ...memo.Implementation) float64 {
+	impl.cost = children[0].GetCost()
+	return impl.cost
+}