[release-22.0] bugfix: allow window functions when possible to push down (#18103) (#18106)

Signed-off-by: Andres Taylor <[email protected]>
Co-authored-by: vitess-bot[bot] <108069721+vitess-bot[bot]@users.noreply.github.com>

Author: vitess-bot[bot]

go/vt/vtgate/planbuilder/operators/plan_query.go

@@ -82,26 +82,36 @@ func PlanQuery(ctx *plancontext.PlanningContext, stmt sqlparser.Statement) (resu
 
 // checkSingleRouteError checks if the query has a NotSingleRouteErr and more than one route, and returns an error if it does
 func checkSingleRouteError(ctx *plancontext.PlanningContext, op Operator) error {
-	if ctx.SemTable.NotSingleRouteErr == nil {
+	if ctx.SemTable.NotSingleRouteErr == nil && ctx.SemTable.NotSingleShardErr == nil {
 		return nil
 	}
+	err := ctx.SemTable.NotSingleRouteErr
+	if err == nil {
+		err = ctx.SemTable.NotSingleShardErr
+	}
 	routes := 0
+	var singleShard bool
 	visitF := func(op Operator, _ semantics.TableSet, _ bool) (Operator, *ApplyResult) {
-		switch op.(type) {
+		switch op := op.(type) {
 		case *Route:
+
 			routes++
+			singleShard = op.IsSingleShard()
 		}
 		return op, NoRewrite
 	}
 
 	// we'll walk the tree and count the number of routes
 	TopDown(op, TableID, visitF, stopAtRoute)
 
-	if routes <= 1 {
-		return nil
+	if routes > 1 {
+		return err
 	}
 
-	return ctx.SemTable.NotSingleRouteErr
+	if ctx.SemTable.NotSingleShardErr != nil && !singleShard {
+		return ctx.SemTable.NotSingleShardErr
+	}
+	return nil
 }
 
 func PanicHandler(err *error) {

go/vt/vtgate/planbuilder/testdata/select_cases.json

@@ -6391,6 +6391,33 @@
     },
     "skip_e2e": true
   },
+  {
+    "comment": "Over clause in a equal unique query going to a single shard",
+    "query": "SELECT Sum(id) OVER ( PARTITION BY name  ) from user where id = 5",
+    "plan": {
+      "Type": "Passthrough",
+      "QueryType": "SELECT",
+      "Original": "SELECT Sum(id) OVER ( PARTITION BY name  ) from user where id = 5",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select sum(id) over ( partition by `name`) from `user` where 1 != 1",
+        "Query": "select sum(id) over ( partition by `name`) from `user` where id = 5",
+        "Table": "`user`",
+        "Values": [
+          "5"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
   {
     "comment": "join with derived table with alias and join condition - merge into route",
     "query": "select 1 from user join (select id as uid from user) as t where t.uid = user.id",

go/vt/vtgate/semantics/analyzer.go

@@ -44,6 +44,7 @@ type analyzer struct {
 	inProjection int
 
 	notSingleRouteErr       error
+	notSingleShardErr       error
 	unshardedErr            error
 	warning                 string
 	canShortcut             bool
@@ -144,6 +145,7 @@ func (a *analyzer) newSemTable(
 			Collation:                 coll,
 			ExprTypes:                 map[sqlparser.Expr]evalengine.Type{},
 			NotSingleRouteErr:         a.notSingleRouteErr,
+			NotSingleShardErr:         a.notSingleShardErr,
 			NotUnshardedErr:           a.unshardedErr,
 			Recursive:                 ExprDependencies{},
 			Direct:                    ExprDependencies{},
@@ -177,6 +179,7 @@ func (a *analyzer) newSemTable(
 		DMLTargets:                a.binder.targets,
 		NotSingleRouteErr:         a.notSingleRouteErr,
 		NotUnshardedErr:           a.unshardedErr,
+		NotSingleShardErr:         a.notSingleShardErr,
 		Warning:                   a.warning,
 		Comments:                  comments,
 		ExprEqualities:            NewTransitiveClosures(),
@@ -198,6 +201,8 @@ func (a *analyzer) setError(err error) {
 		a.notSingleRouteErr = err.Inner
 	case ShardedError:
 		a.unshardedErr = err.Inner
+	case NotSingleShardError:
+		a.notSingleShardErr = err.Inner
 	default:
 		if a.inProjection > 0 && vterrors.ErrState(err) == vterrors.NonUniqError {
 			a.notSingleRouteErr = err
@@ -491,26 +496,37 @@ func (a *analyzer) getError() error {
 	return a.err
 }
 
-// NotSingleRouteErr is used to mark an error as something that should only be returned
-// if the planner fails to merge everything down to a single route
-type NotSingleRouteErr struct {
-	Inner error
-}
+type (
+	// NotSingleRouteErr is used to mark an error as something that should only be returned
+	// if the planner fails to merge everything down to a single route
+	NotSingleRouteErr struct {
+		Inner error
+	}
+	// ShardedError is used to mark an error as something that should only be returned
+	// if the query is not unsharded
+	ShardedError struct {
+		Inner error
+	}
+	// NotSingleShardError is used to mark an error as something that should only be returned
+	// if the query fails to be planned into a single shard query
+	NotSingleShardError struct {
+		Inner error
+	}
+)
 
 func (p NotSingleRouteErr) Error() string {
 	return p.Inner.Error()
 }
+func (p NotSingleRouteErr) Unwrap() error { return p.Inner }
 
-// ShardedError is used to mark an error as something that should only be returned
-// if the query is not unsharded
-type ShardedError struct {
-	Inner error
+func (p ShardedError) Error() string {
+	return p.Inner.Error()
 }
-
 func (p ShardedError) Unwrap() error {
 	return p.Inner
 }
 
-func (p ShardedError) Error() string {
+func (p NotSingleShardError) Error() string {
 	return p.Inner.Error()
 }
+func (p NotSingleShardError) Unwrap() error { return p.Inner }

go/vt/vtgate/semantics/binder.go

@@ -17,6 +17,7 @@ limitations under the License.
 package semantics
 
 import (
+	"errors"
 	"strings"
 
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -303,22 +304,11 @@ func (b *binder) resolveColumn(colName *sqlparser.ColName, current *scope, allow
 	return dependency{}, ShardedError{ColumnNotFoundError{Column: colName, Table: tableName}}
 }
 
-func isColumnNotFound(err error) bool {
-	switch err := err.(type) {
-	case ColumnNotFoundError:
-		return true
-	case ShardedError:
-		return isColumnNotFound(err.Inner)
-	default:
-		return false
-	}
-}
-
 func (b *binder) resolveColumnInHaving(colName *sqlparser.ColName, current *scope, allowMulti bool) (dependency, error) {
 	if current.inHavingAggr {
 		// when inside an aggregation, we'll search the FROM clause before the SELECT expressions
 		deps, err := b.resolveColumn(colName, current.parent, allowMulti, true)
-		if deps.direct.NotEmpty() || (err != nil && !isColumnNotFound(err)) {
+		if deps.direct.NotEmpty() || (err != nil && !errors.As(err, &ColumnNotFoundError{})) {
 			return deps, err
 		}
 	}
@@ -354,7 +344,7 @@ func (b *binder) resolveColumnInHaving(colName *sqlparser.ColName, current *scop
 
 	if !current.inHavingAggr && sel.GroupBy == nil {
 		// if we are not inside an aggregation, and there is no GROUP BY, we consider the FROM clause before failing
-		if deps.direct.NotEmpty() || (err != nil && !isColumnNotFound(err)) {
+		if deps.direct.NotEmpty() || (err != nil && !errors.As(err, &ColumnNotFoundError{})) {
 			return deps, err
 		}
 	}
@@ -429,7 +419,7 @@ func (b *binder) resolveColInGroupBy(
 		return dependency{}, err
 	}
 	if dependencies.empty() {
-		if isColumnNotFound(firstErr) {
+		if errors.As(firstErr, &ColumnNotFoundError{}) {
 			return dependency{}, &ColumnNotFoundClauseError{Column: colName.Name.String(), Clause: "group statement"}
 		}
 		return deps, firstErr

go/vt/vtgate/semantics/check_invalid.go

@@ -54,7 +54,7 @@ func (a *analyzer) checkForInvalidConstructs(cursor *sqlparser.Cursor) error {
 		}
 	case *sqlparser.OverClause:
 		if !a.singleUnshardedKeyspace {
-			return ShardedError{Inner: &UnsupportedConstruct{errString: "OVER CLAUSE with sharded keyspace"}}
+			return NotSingleShardError{Inner: &UnsupportedConstruct{errString: "OVER CLAUSE with sharded keyspace"}}
 		}
 	}
 

go/vt/vtgate/semantics/semantic_table.go

@@ -122,6 +122,10 @@ type (
 		// MySQL engine to handle errors appropriately.
 		NotUnshardedErr error
 
+		// If there are constructs in this query that we know we only support if we can push them down
+		// unbroken to mysql, this field will contain an error that is produced when we fail to do so.
+		NotSingleShardErr error
+
 		// Recursive contains dependencies from the expression to the actual tables
 		// in the query (excluding derived tables). For columns in derived tables,
 		// this map holds the accumulated dependencies for the column expression.