Merge pull request #8961 from planetscale/gen4-orderby-proj

Gen4: support for ordering using derived table columns

go/vt/vtgate/planbuilder/abstract/queryprojection.go

@@ -253,14 +253,37 @@ func (qp *QueryProjection) getNonAggrExprNotMatchingGroupByExprs() sqlparser.Sel
 	return nil
 }
 
-// getSimplifiedExpr takes an expression used in ORDER BY or GROUP BY, which can reference both aliased columns and
-// column offsets, and returns an expression that is simpler to evaluate
+// getSimplifiedExpr takes an expression used in ORDER BY or GROUP BY, and returns an expression that is simpler to evaluate
 func (qp *QueryProjection) getSimplifiedExpr(e sqlparser.Expr, semTable *semantics.SemTable) (expr sqlparser.Expr, weightStrExpr sqlparser.Expr, err error) {
 	// If the ORDER BY is against a column alias, we need to remember the expression
 	// behind the alias. The weightstring(.) calls needs to be done against that expression and not the alias.
 	// Eg - select music.foo as bar, weightstring(music.foo) from music order by bar
+
 	colExpr, isColName := e.(*sqlparser.ColName)
-	if isColName && colExpr.Qualifier.IsEmpty() {
+	if !isColName {
+		return e, e, nil
+	}
+
+	if sqlparser.IsNull(e) {
+		return e, nil, nil
+	}
+
+	tblInfo, err := semTable.TableInfoForExpr(e)
+	if err != nil && err != semantics.ErrMultipleTables {
+		// we can live with ErrMultipleTables and just ignore it. anything else should fail this method
+		return nil, nil, err
+	}
+	if tblInfo != nil {
+		if dTablInfo, ok := tblInfo.(*semantics.DerivedTable); ok {
+			weightStrExpr, err = semantics.RewriteDerivedExpression(colExpr, dTablInfo)
+			if err != nil {
+				return nil, nil, err
+			}
+			return e, weightStrExpr, nil
+		}
+	}
+
+	if colExpr.Qualifier.IsEmpty() {
 		for _, selectExpr := range qp.SelectExprs {
 			aliasedExpr, isAliasedExpr := selectExpr.Col.(*sqlparser.AliasedExpr)
 			if !isAliasedExpr {
@@ -273,10 +296,6 @@ func (qp *QueryProjection) getSimplifiedExpr(e sqlparser.Expr, semTable *semanti
 		}
 	}
 
-	if sqlparser.IsNull(e) {
-		return e, nil, nil
-	}
-
 	return e, e, nil
 }
 

go/vt/vtgate/planbuilder/horizon_planning.go

@@ -274,10 +274,9 @@ func pushProjection(expr *sqlparser.AliasedExpr, plan logicalPlan, semTable *sem
 }
 
 func rewriteProjectionOfDerivedTable(expr *sqlparser.AliasedExpr, semTable *semantics.SemTable) error {
-	var err error
-	ti, _ := semTable.TableInfoForExpr(expr.Expr)
-	if ti == nil {
-		return nil
+	ti, err := semTable.TableInfoForExpr(expr.Expr)
+	if err != nil && err != semantics.ErrMultipleTables {
+		return err
 	}
 	_, isDerivedTable := ti.(*semantics.DerivedTable)
 	if isDerivedTable {
@@ -630,7 +629,7 @@ func (hp *horizonPlanning) planOrderBy(ctx *planningContext, orderExprs []abstra
 		plan.input = newUnderlyingPlan
 		return plan, nil
 	case *simpleProjection:
-		return nil, semantics.Gen4NotSupportedF("unsupported: ordering on derived table query")
+		return hp.createMemorySortPlan(ctx, plan, orderExprs)
 	case *vindexFunc:
 		return nil, semantics.Gen4NotSupportedF("unsupported: ordering on vindex func")
 	default:

go/vt/vtgate/planbuilder/testdata/memory_sort_cases.txt

@@ -313,6 +313,57 @@ Gen4 error: Expression of SELECT list is not in GROUP BY clause and contains non
     ]
   }
 }
+{
+  "QueryType": "SELECT",
+  "Original": "select id from (select user.id, user.col from user join user_extra) as t order by id",
+  "Instructions": {
+    "OperatorType": "Sort",
+    "Variant": "Memory",
+    "OrderBy": "(0|1) ASC",
+    "ResultColumns": 1,
+    "Inputs": [
+      {
+        "OperatorType": "SimpleProjection",
+        "Columns": [
+          0,
+          2
+        ],
+        "Inputs": [
+          {
+            "OperatorType": "Join",
+            "Variant": "Join",
+            "JoinColumnIndexes": "-1,-2,-3",
+            "TableName": "`user`_user_extra",
+            "Inputs": [
+              {
+                "OperatorType": "Route",
+                "Variant": "SelectScatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select `user`.id, `user`.col, weight_string(`user`.id) from `user` where 1 != 1",
+                "Query": "select `user`.id, `user`.col, weight_string(`user`.id) from `user`",
+                "Table": "`user`"
+              },
+              {
+                "OperatorType": "Route",
+                "Variant": "SelectScatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select 1 from user_extra where 1 != 1",
+                "Query": "select 1 from user_extra",
+                "Table": "user_extra"
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+}
 
 # order by on a cross-shard query. Note: this happens only when an order by column is from the second table
 "select user.col1 as a, user.col2 b, music.col3 c from user, music where user.id = music.id and user.id = 1 order by c"

go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt

@@ -30,6 +30,7 @@ Gen4 error: unsupported: '*' expression in cross-shard query
 # order by rand on a cross-shard subquery
 "select id from (select user.id, user.col from user join user_extra) as t order by rand()"
 "unsupported: memory sort: complex order by expression: rand()"
+Gen4 error: unsupported: in scatter query: complex order by expression: rand()
 
 # natural join
 "select * from user natural join user_extra"

go/vt/vtgate/semantics/analyzer.go

@@ -216,16 +216,18 @@ func isParentSelect(cursor *sqlparser.Cursor) bool {
 
 type originable interface {
 	tableSetFor(t *sqlparser.AliasedTableExpr) TableSet
-	depsForExpr(expr sqlparser.Expr) (TableSet, *querypb.Type)
+	depsForExpr(expr sqlparser.Expr) (direct, recursive TableSet, typ *querypb.Type)
 }
 
-func (a *analyzer) depsForExpr(expr sqlparser.Expr) (TableSet, *querypb.Type) {
-	ts := a.binder.recursive.Dependencies(expr)
+func (a *analyzer) depsForExpr(expr sqlparser.Expr) (direct, recursive TableSet, typ *querypb.Type) {
+	recursive = a.binder.recursive.Dependencies(expr)
+	direct = a.binder.direct.Dependencies(expr)
 	qt, isFound := a.typer.exprTypes[expr]
 	if !isFound {
-		return ts, nil
+		return
 	}
-	return ts, &qt
+	typ = &qt
+	return
 }
 
 func (a *analyzer) analyze(statement sqlparser.Statement) error {

go/vt/vtgate/semantics/analyzer_test.go

@@ -927,6 +927,65 @@ func TestScopingWDerivedTables(t *testing.T) {
 	}
 }
 
+func TestDerivedTablesOrderClause(t *testing.T) {
+	queries := []struct {
+		query                string
+		recursiveExpectation TableSet
+		expectation          TableSet
+	}{{
+		query:                "select 1 from (select id from user) as t order by id",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select id from (select id from user) as t order by id",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select id from (select id from user) as t order by t.id",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select id as foo from (select id from user) as t order by foo",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select bar from (select id as bar from user) as t order by bar",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select bar as foo from (select id as bar from user) as t order by bar",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select bar as foo from (select id as bar from user) as t order by foo",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select bar as foo from (select id as bar, oo from user) as t order by oo",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}, {
+		query:                "select bar as foo from (select id, oo from user) as t(bar,oo) order by bar",
+		recursiveExpectation: T1,
+		expectation:          T2,
+	}}
+	si := &FakeSI{Tables: map[string]*vindexes.Table{"t": {Name: sqlparser.NewTableIdent("t")}}}
+	for _, query := range queries {
+		t.Run(query.query, func(t *testing.T) {
+			parse, err := sqlparser.Parse(query.query)
+			require.NoError(t, err)
+
+			st, err := Analyze(parse.(sqlparser.SelectStatement), "user", si)
+			require.NoError(t, err)
+
+			sel := parse.(*sqlparser.Select)
+			assert.Equal(t, query.recursiveExpectation, st.RecursiveDeps(sel.OrderBy[0].Expr), "RecursiveDeps")
+			assert.Equal(t, query.expectation, st.DirectDeps(sel.OrderBy[0].Expr), "DirectDeps")
+
+		})
+	}
+}
+
 func TestScopingWComplexDerivedTables(t *testing.T) {
 	queries := []struct {
 		query            string

go/vt/vtgate/semantics/derived_table.go

@@ -23,6 +23,7 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
 )
 
+// DerivedTable contains the information about the projection, tables involved in derived table.
 type DerivedTable struct {
 	tableName   string
 	ASTNode     *sqlparser.AliasedTableExpr
@@ -63,27 +64,21 @@ func createDerivedTableForExpressions(expressions sqlparser.SelectExprs, cols sq
 
 // Dependencies implements the TableInfo interface
 func (dt *DerivedTable) dependencies(colName string, org originable) (dependencies, error) {
+	directDeps := org.tableSetFor(dt.ASTNode)
 	for i, name := range dt.columnNames {
 		if name != colName {
 			continue
 		}
-		recursiveDeps, qt := org.depsForExpr(dt.cols[i])
+		_, recursiveDeps, qt := org.depsForExpr(dt.cols[i])
 
-		directDeps := recursiveDeps
-		if dt.ASTNode != nil {
-			directDeps = org.tableSetFor(dt.ASTNode)
-		}
 		return createCertain(directDeps, recursiveDeps, qt), nil
 	}
 
 	if !dt.hasStar() {
 		return &nothing{}, nil
 	}
 
-	recursive := dt.tables
-	direct := org.tableSetFor(dt.ASTNode)
-
-	return createUncertain(direct, recursive), nil
+	return createUncertain(directDeps, dt.tables), nil
 }
 
 // IsInfSchema implements the TableInfo interface
@@ -99,6 +94,7 @@ func (dt *DerivedTable) authoritative() bool {
 	return true
 }
 
+// Name implements the TableInfo interface
 func (dt *DerivedTable) Name() (sqlparser.TableName, error) {
 	return dt.ASTNode.TableName()
 }

go/vt/vtgate/semantics/scoper.go

@@ -83,13 +83,13 @@ func (s *scoper) down(cursor *sqlparser.Cursor) error {
 			break
 		}
 
-		// adding a VTableInfo for each SELECT, so it can be used by GROUP BY, HAVING, ORDER BY
-		// the VTableInfo we are creating here should not be confused with derived tables' VTableInfo
+		// adding a vTableInfo for each SELECT, so it can be used by GROUP BY, HAVING, ORDER BY
+		// the vTableInfo we are creating here should not be confused with derived tables' vTableInfo
 		wScope, exists := s.wScope[sel]
 		if !exists {
 			break
 		}
-		wScope.tables = append(wScope.tables, createVTableInfoForExpressions(node, s.currentScope().tables, s.org))
+		wScope.tables = []TableInfo{createVTableInfoForExpressions(node, s.currentScope().tables, s.org)}
 	case sqlparser.OrderBy:
 		err := s.createSpecialScopePostProjection(cursor.Parent())
 		if err != nil {
@@ -171,10 +171,9 @@ func (s *scoper) createSpecialScopePostProjection(parent sqlparser.SQLNode) erro
 		// so before walking the rest of the tree, we change the scope to match this behaviour
 		incomingScope := s.currentScope()
 		nScope := newScope(incomingScope)
-		s.push(nScope)
-		wScope := s.wScope[parent]
-		nScope.tables = append(nScope.tables, wScope.tables...)
+		nScope.tables = s.wScope[parent].tables
 		nScope.selectStmt = incomingScope.selectStmt
+		s.push(nScope)
 
 		if s.rScope[parent] != incomingScope {
 			return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "BUG: scope counts did not match")

go/vt/vtgate/semantics/semantic_state.go

@@ -129,7 +129,11 @@ func (st *SemTable) CopyDependencies(from, to sqlparser.Expr) {
 
 // NewSemTable creates a new empty SemTable
 func NewSemTable() *SemTable {
-	return &SemTable{Recursive: map[sqlparser.Expr]TableSet{}, ColumnEqualities: map[columnName][]sqlparser.Expr{}}
+	return &SemTable{
+		Recursive:        map[sqlparser.Expr]TableSet{},
+		Direct:           map[sqlparser.Expr]TableSet{},
+		ColumnEqualities: map[columnName][]sqlparser.Expr{},
+	}
 }
 
 // TableSetFor returns the bitmask for this particular table

go/vt/vtgate/semantics/vtable.go

@@ -42,20 +42,7 @@ func (v *vTableInfo) dependencies(colName string, org originable) (dependencies,
 		if name != colName {
 			continue
 		}
-		recursiveDeps, qt := org.depsForExpr(v.cols[i])
-
-		var directDeps TableSet
-		/*
-				If we find a match, it means the query looks something like:
-				SELECT 1 as x FROM t1 ORDER BY/GROUP BY x - d/r: 0/0
-				SELECT t1.x as x FROM t1 ORDER BY/GROUP BY x - d/r: 0/1
-				SELECT x FROM t1 ORDER BY/GROUP BY x - d/r: 1/1
-
-			    Now, after figuring out the recursive deps
-		*/
-		if recursiveDeps.NumberOfTables() > 0 {
-			directDeps = recursiveDeps
-		}
+		directDeps, recursiveDeps, qt := org.depsForExpr(v.cols[i])
 
 		newDeps := createCertain(directDeps, recursiveDeps, qt)
 		deps, err = deps.merge(newDeps)