planbuilder: add support for window functions

Signed-off-by: crangavajha1 <[email protected]>

Author: c-r-dev

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -91,6 +91,8 @@ func recursiveTransform(ctx *plancontext.PlanningContext, op operators.Operator)
 		return transformDMLWithInput(ctx, op)
 	case *operators.RecurseCTE:
 		return transformRecurseCTE(ctx, op)
+	case *operators.Window:
+		return transformWindow(ctx, op)
 	case *operators.PercentBasedMirror:
 		return transformPercentBasedMirror(ctx, op)
 	}
@@ -1059,3 +1061,23 @@ func generateQuery(statement sqlparser.Statement) string {
 	statement.Format(buf)
 	return buf.String()
 }
+
+func isSingleShard(prim engine.Primitive) bool {
+	switch p := prim.(type) {
+	case *engine.Route:
+		return p.Opcode.IsSingleShard() || p.Opcode == engine.ByDestination
+	case *engine.PlanSwitcher:
+		return isSingleShard(p.Optimized)
+	case *engine.Filter:
+		return isSingleShard(p.Input)
+	case *engine.Limit:
+		return isSingleShard(p.Input)
+	case *engine.MemorySort:
+		return isSingleShard(p.Input)
+	case *engine.Projection:
+		return isSingleShard(p.Input)
+	case *engine.SimpleProjection:
+		return isSingleShard(p.Input)
+	}
+	return false
+}

go/vt/vtgate/planbuilder/operators/horizon_expanding.go

@@ -94,6 +94,14 @@ func expandSelectHorizon(ctx *plancontext.PlanningContext, horizon *Horizon, sel
 		extracted = append(extracted, "Projection")
 	}
 
+	if qp.HasWindow {
+		// Window functions are evaluated after aggregation but before ordering and limit.
+		// We wrap the current operator (which is either a Projection or Aggregation)
+		// with the Window operator to handle these calculations.
+		op = newWindow(op, qp)
+		extracted = append(extracted, "Window")
+	}
+
 	if qp.NeedsDistinct() {
 		op = newDistinct(op, qp, true)
 		extracted = append(extracted, "Distinct")

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -292,6 +292,7 @@ func pushOrExpandHorizon(ctx *plancontext.PlanningContext, in *Horizon) (Operato
 		!hasHaving &&
 		!needsOrdering &&
 		!qp.NeedsAggregation() &&
+		!qp.HasWindow &&
 		!isDistinctAST(in.selectStatement()) &&
 		in.selectStatement().GetLimit() == nil
 
@@ -308,6 +309,8 @@ func pushOrExpandHorizon(ctx *plancontext.PlanningContext, in *Horizon) (Operato
 		debugNoRewrite("horizon push blocked: query has ORDER BY")
 	} else if qp.NeedsAggregation() {
 		debugNoRewrite("horizon push blocked: query needs aggregation")
+	} else if qp.HasWindow {
+		debugNoRewrite("horizon push blocked: query has window functions")
 	} else if isDistinctAST(in.selectStatement()) {
 		debugNoRewrite("horizon push blocked: query has DISTINCT")
 	} else if in.selectStatement().GetLimit() != nil {

go/vt/vtgate/planbuilder/operators/queryprojection.go

@@ -42,6 +42,7 @@ type (
 		// If you change the contents here, please update the toString() method
 		SelectExprs  []SelectExpr
 		HasAggr      bool
+		HasWindow    bool
 		Distinct     bool
 		WithRollup   bool
 		groupByExprs []GroupBy
@@ -169,6 +170,14 @@ func createQPFromSelect(ctx *plancontext.PlanningContext, sel *sqlparser.Select)
 	if !qp.HasAggr && sel.Having != nil {
 		qp.HasAggr = ctx.ContainsAggr(sel.Having.Expr)
 	}
+	if !qp.HasWindow {
+		for _, order := range sel.OrderBy {
+			if ctx.ContainsWindowFunc(order.Expr) {
+				qp.HasWindow = true
+				break
+			}
+		}
+	}
 	qp.calculateDistinct(ctx)
 
 	return qp
@@ -185,6 +194,9 @@ func (qp *QueryProjection) addSelectExpressions(ctx *plancontext.PlanningContext
 				col.Aggr = true
 				qp.HasAggr = true
 			}
+			if ctx.ContainsWindowFunc(selExp.Expr) {
+				qp.HasWindow = true
+			}
 
 			qp.SelectExprs = append(qp.SelectExprs, col)
 		case *sqlparser.StarExpr:
@@ -414,6 +426,12 @@ func (qp *QueryProjection) AggregationExpressions(ctx *plancontext.PlanningConte
 			if err != nil {
 				panic(err)
 			}
+
+			if ctx.ContainsWindowFunc(selectExpr.Col) {
+				sqlparser.CopyOnRewrite(aliasedExpr.Expr, qp.extractAggr(ctx, aliasedExpr, addAggr, makeComplex), nil, nil)
+				continue
+			}
+
 			if !qp.isExprInGroupByExprs(ctx, getExpr) {
 				aggr := createNonGroupingAggr(aliasedExpr)
 				out = append(out, aggr)
@@ -441,6 +459,9 @@ func (qp *QueryProjection) extractAggr(
 			return true
 		}
 		if aggr, isAggr := node.(sqlparser.AggrFunc); isAggr {
+			if wf, ok := node.(sqlparser.WindowFunc); ok && wf.GetOverClause() != nil {
+				return true
+			}
 			ae := aeWrap(aggr)
 			if aggr == aliasedExpr.Expr {
 				ae = aliasedExpr

go/vt/vtgate/planbuilder/operators/subquery_planning.go

@@ -43,6 +43,12 @@ func isMergeable(ctx *plancontext.PlanningContext, query sqlparser.TableStatemen
 
 	switch node := query.(type) {
 	case *sqlparser.Select:
+		// Window functions cannot be merged into the outer scope because they operate on a result set
+		// and require specific partitioning and ordering semantics that would be lost in a merge
+		if ctx.ContainsWindowFunc(node) {
+			return false
+		}
+
 		if node.GroupBy != nil && len(node.GroupBy.Exprs) > 0 {
 			// iff we are grouping, we need to check that we can perform the grouping inside a single shard, and we check that
 			// by checking that one of the grouping expressions used is a unique single column vindex.

go/vt/vtgate/planbuilder/operators/window.go

@@ -0,0 +1,71 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package operators
+
+import (
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+)
+
+type Window struct {
+	unaryOperator
+	QP *QueryProjection
+}
+
+func newWindow(source Operator, qp *QueryProjection) *Window {
+	return &Window{
+		unaryOperator: newUnaryOp(source),
+		QP:            qp,
+	}
+}
+
+func (w *Window) Clone(inputs []Operator) Operator {
+	return newWindow(inputs[0], w.QP)
+}
+
+func (w *Window) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Expr) Operator {
+	w.Source = w.Source.AddPredicate(ctx, expr)
+	return w
+}
+
+func (w *Window) AddColumn(ctx *plancontext.PlanningContext, reuseExisting bool, addToGroupBy bool, expr *sqlparser.AliasedExpr) int {
+	return w.Source.AddColumn(ctx, reuseExisting, addToGroupBy, expr)
+}
+
+func (w *Window) AddWSColumn(ctx *plancontext.PlanningContext, offset int, underRoute bool) int {
+	return w.Source.AddWSColumn(ctx, offset, underRoute)
+}
+
+func (w *Window) FindCol(ctx *plancontext.PlanningContext, expr sqlparser.Expr, underRoute bool) int {
+	return w.Source.FindCol(ctx, expr, underRoute)
+}
+
+func (w *Window) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr {
+	return w.Source.GetColumns(ctx)
+}
+
+func (w *Window) GetSelectExprs(ctx *plancontext.PlanningContext) []sqlparser.SelectExpr {
+	return w.Source.GetSelectExprs(ctx)
+}
+
+func (w *Window) ShortDescription() string {
+	return "Window"
+}
+
+func (w *Window) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
+	return w.Source.GetOrdering(ctx)
+}

go/vt/vtgate/planbuilder/plan_test.go

@@ -98,6 +98,7 @@ func (s *planTestSuite) TestPlan() {
 	s.testFile("filter_cases.json", vw, false)
 	s.testFile("postprocess_cases.json", vw, false)
 	s.testFile("select_cases.json", vw, false)
+	s.testFile("window_function_cases.json", vw, false)
 	s.testFile("symtab_cases.json", vw, false)
 	s.testFile("unsupported_cases.json", vw, false)
 	s.testFile("unknown_schema_cases.json", vw, false)

go/vt/vtgate/planbuilder/plancontext/planning_context.go

@@ -277,12 +277,15 @@ func (ctx *PlanningContext) IsAggr(e sqlparser.SQLNode) bool {
 
 func (ctx *PlanningContext) ContainsAggr(e sqlparser.SQLNode) (hasAggr bool) {
 	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
-		switch node.(type) {
+		switch node := node.(type) {
 		case *sqlparser.Offset:
 			// offsets here indicate that a possible aggregation has already been handled by an input,
 			// so we don't need to worry about aggregation in the original
 			return false, nil
 		case sqlparser.AggrFunc:
+			if wf, ok := node.(sqlparser.WindowFunc); ok && wf.GetOverClause() != nil {
+				return true, nil
+			}
 			hasAggr = true
 			return false, io.EOF
 		case *sqlparser.Subquery:
@@ -299,6 +302,26 @@ func (ctx *PlanningContext) ContainsAggr(e sqlparser.SQLNode) (hasAggr bool) {
 	return
 }
 
+func (ctx *PlanningContext) ContainsWindowFunc(e sqlparser.SQLNode) (hasWindow bool) {
+	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		switch node := node.(type) {
+		case *sqlparser.Offset:
+			// offsets here indicate that a possible window function has already been handled by an input,
+			// so we don't need to worry about it in the original
+			return false, nil
+		case sqlparser.WindowFunc:
+			if node.GetOverClause() != nil {
+				hasWindow = true
+				return false, io.EOF
+			}
+		case *sqlparser.Subquery:
+			return false, nil
+		}
+		return true, nil
+	}, e)
+	return
+}
+
 func (ctx *PlanningContext) IsMirrored() bool {
 	return ctx.isMirrored
 }