Merge pull request #7775 from planetscale/gen4-thing

[Gen4] Initial Horizon Planning

go/mysql/sql_error.go

@@ -170,6 +170,7 @@ var stateToMysqlCode = map[vterrors.State]struct {
 	vterrors.NotSupportedYet:              {num: ERNotSupportedYet, state: SSClientError},
 	vterrors.ForbidSchemaChange:           {num: ERForbidSchemaChange, state: SSUnknownSQLState},
 	vterrors.NetPacketTooLarge:            {num: ERNetPacketTooLarge, state: SSNetError},
+	vterrors.NonUniqError:                 {num: ERNonUniq, state: SSConstraintViolation},
 	vterrors.NonUniqTable:                 {num: ERNonUniqTable, state: SSClientError},
 	vterrors.QueryInterrupted:             {num: ERQueryInterrupted, state: SSQueryInterrupted},
 	vterrors.SPDoesNotExist:               {num: ERSPDoesNotExist, state: SSClientError},

go/vt/vterrors/state.go

@@ -30,6 +30,7 @@ const (
 	EmptyQuery
 	ForbidSchemaChange
 	IncorrectGlobalLocalVar
+	NonUniqError
 	NonUniqTable
 	SyntaxError
 	WrongGroupField

go/vt/vtgate/engine/route.go

@@ -739,6 +739,9 @@ func (route *Route) description() PrimitiveDescription {
 	if orderBy != "" {
 		other["OrderBy"] = orderBy
 	}
+	if route.TruncateColumnCount > 0 {
+		other["ResultColumns"] = route.TruncateColumnCount
+	}
 
 	return PrimitiveDescription{
 		OperatorType:      "Route",

go/vt/vtgate/planbuilder/querygraph_test.go

@@ -22,6 +22,10 @@ import (
 	"strings"
 	"testing"
 
+	"vitess.io/vitess/go/vt/key"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/vtgate/vindexes"
+
 	"github.com/stretchr/testify/assert"
 
 	"vitess.io/vitess/go/test/utils"
@@ -119,7 +123,7 @@ func TestQueryGraph(t *testing.T) {
 		t.Run(fmt.Sprintf("%d %s", i, sql), func(t *testing.T) {
 			tree, err := sqlparser.Parse(sql)
 			require.NoError(t, err)
-			semTable, err := semantics.Analyse(tree)
+			semTable, err := semantics.Analyse(tree, &fakeSI{})
 			require.NoError(t, err)
 			qgraph, err := createQGFromSelect(tree.(*sqlparser.Select), semTable)
 			require.NoError(t, err)
@@ -132,7 +136,7 @@ func TestQueryGraph(t *testing.T) {
 func TestString(t *testing.T) {
 	tree, err := sqlparser.Parse("select * from a,b join c on b.id = c.id where a.id = b.id and b.col IN (select 42) and func() = 'foo'")
 	require.NoError(t, err)
-	semTable, err := semantics.Analyse(tree)
+	semTable, err := semantics.Analyse(tree, &fakeSI{})
 	require.NoError(t, err)
 	qgraph, err := createQGFromSelect(tree.(*sqlparser.Select), semTable)
 	require.NoError(t, err)
@@ -232,3 +236,13 @@ func (qg *queryGraph) noDepsString() string {
 	}
 	return fmt.Sprintf("\nForAll: %s", sqlparser.String(qg.noDeps))
 }
+
+var _ semantics.SchemaInformation = (*fakeSI)(nil)
+
+type fakeSI struct {
+	tables map[string]*vindexes.Table
+}
+
+func (s *fakeSI) FindTableOrVindex(tablename sqlparser.TableName) (*vindexes.Table, vindexes.Vindex, string, topodatapb.TabletType, key.Destination, error) {
+	return s.tables[sqlparser.String(tablename)], nil, "", 0, nil, nil
+}

go/vt/vtgate/planbuilder/queryprojection.go

@@ -0,0 +1,84 @@
+/*
+Copyright 2021 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 planbuilder
+
+import (
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vterrors"
+	"vitess.io/vitess/go/vt/vtgate/semantics"
+)
+
+type queryProjection struct {
+	selectExprs             []*sqlparser.AliasedExpr
+	aggrExprs               []*sqlparser.AliasedExpr
+	groupOrderingCommonExpr map[sqlparser.Expr]*sqlparser.Order
+
+	orderExprs sqlparser.OrderBy
+
+	// orderExprColMap keeps a map between the Order object and the offset into the select expressions list
+	orderExprColMap map[*sqlparser.Order]int
+}
+
+func newQueryProjection() *queryProjection {
+	return &queryProjection{
+		groupOrderingCommonExpr: map[sqlparser.Expr]*sqlparser.Order{},
+		orderExprColMap:         map[*sqlparser.Order]int{},
+	}
+}
+
+func createQPFromSelect(sel *sqlparser.Select) (*queryProjection, error) {
+	qp := newQueryProjection()
+
+	for _, selExp := range sel.SelectExprs {
+		exp, ok := selExp.(*sqlparser.AliasedExpr)
+		if !ok {
+			return nil, semantics.Gen4NotSupportedF("%T in select list", selExp)
+		}
+		fExpr, ok := exp.Expr.(*sqlparser.FuncExpr)
+		if ok && fExpr.IsAggregate() {
+			if len(fExpr.Exprs) != 1 {
+				return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.SyntaxError, "multiple arguments inside the function '%s'", sqlparser.String(fExpr))
+			}
+			qp.aggrExprs = append(qp.aggrExprs, exp)
+			continue
+		}
+		if nodeHasAggregates(exp.Expr) {
+			return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: in scatter query: complex aggregate expression")
+		}
+		qp.selectExprs = append(qp.selectExprs, exp)
+	}
+
+	qp.orderExprs = sel.OrderBy
+
+	allExpr := append(qp.selectExprs, qp.aggrExprs...)
+	for _, order := range sel.OrderBy {
+		for offset, expr := range allExpr {
+			if sqlparser.EqualsExpr(order.Expr, expr.Expr) {
+				qp.orderExprColMap[order] = offset
+				break
+			}
+			// TODO: handle alias and column offset
+		}
+	}
+
+	if sel.GroupBy == nil || sel.OrderBy == nil {
+		return qp, nil
+	}
+
+	return qp, nil
+}

go/vt/vtgate/planbuilder/route_planning.go

@@ -46,7 +46,7 @@ func gen4Planner(_ string) func(sqlparser.Statement, sqlparser.BindVars, Context
 }
 
 func newBuildSelectPlan(sel *sqlparser.Select, vschema ContextVSchema) (engine.Primitive, error) {
-	semTable, err := semantics.Analyse(sel) // TODO no nil no
+	semTable, err := semantics.Analyse(sel, vschema)
 	if err != nil {
 		return nil, err
 	}
@@ -74,7 +74,8 @@ func newBuildSelectPlan(sel *sqlparser.Select, vschema ContextVSchema) (engine.P
 		return nil, err
 	}
 
-	if err := planProjections(sel, plan, semTable); err != nil {
+	plan, err = planHorizon(sel, plan, semTable)
+	if err != nil {
 		return nil, err
 	}
 
@@ -112,7 +113,7 @@ func planLimit(limit *sqlparser.Limit, plan logicalPlan) (logicalPlan, error) {
 	return lPlan, nil
 }
 
-func planProjections(sel *sqlparser.Select, plan logicalPlan, semTable *semantics.SemTable) error {
+func planHorizon(sel *sqlparser.Select, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) {
 	rb, ok := plan.(*route)
 	if ok && rb.isSingleShard() {
 		ast := rb.Select.(*sqlparser.Select)
@@ -121,30 +122,40 @@ func planProjections(sel *sqlparser.Select, plan logicalPlan, semTable *semantic
 		ast.OrderBy = sel.OrderBy
 		ast.SelectExprs = sel.SelectExprs
 		ast.Comments = sel.Comments
-	} else {
-		// TODO real horizon planning to be done
-		if sel.Distinct {
-			return semantics.Gen4NotSupportedF("DISTINCT")
+		return plan, nil
+	}
+
+	// TODO real horizon planning to be done
+	if sel.Distinct {
+		return nil, semantics.Gen4NotSupportedF("DISTINCT")
+	}
+	if sel.GroupBy != nil {
+		return nil, semantics.Gen4NotSupportedF("GROUP BY")
+	}
+	qp, err := createQPFromSelect(sel)
+	if err != nil {
+		return nil, err
+	}
+	for _, e := range qp.selectExprs {
+		if _, err := pushProjection(e, plan, semTable); err != nil {
+			return nil, err
 		}
-		if sel.GroupBy != nil {
-			return semantics.Gen4NotSupportedF("GROUP BY")
+	}
+
+	if len(qp.aggrExprs) > 0 {
+		plan, err = planAggregations(qp, plan, semTable)
+		if err != nil {
+			return nil, err
 		}
-		for _, expr := range sel.SelectExprs {
-			switch e := expr.(type) {
-			case *sqlparser.AliasedExpr:
-				if nodeHasAggregates(e.Expr) {
-					return semantics.Gen4NotSupportedF("aggregation [%s]", sqlparser.String(e))
-				}
-				if _, err := pushProjection(e, plan, semTable); err != nil {
-					return err
-				}
-			default:
-				return semantics.Gen4NotSupportedF("%T", e)
-			}
+	}
+	if len(sel.OrderBy) > 0 {
+		plan, err = planOrderBy(qp, plan, semTable)
+		if err != nil {
+			return nil, err
 		}
-
 	}
-	return nil
+
+	return plan, nil
 }
 
 type (