Repository: vitessio/vitess
Pull request #9869 (2022-03-11 08:32:15)

fix: planner panic on derived tables sorting in query builder 
<!--
  Thank you for your contribution to the Vitess project.
  How to contribute: https://vitess.io/docs/contributing/
  Please first make sure there is an open Issue to discuss the feature/fix suggested in this PR.
  If this is a new feature, please mark the Issue as "RFC".
 -->

<!-- if this PR is Work in Progress please create it as a Draft Pull Request -->

## Description
<!-- A few sentences describing the overall goals of the pull request's commits. -->
<!-- If this is a bug fix and you think the fix should be backported, please write so. -->

This PR fixes a panic on table sorting on the final plan query output. It was caused when a query contains a join inside a derived table.

## Related Issue(s)
<!-- List related issues and pull requests. If this PR fixes an issue, please add it using Fixes #????  -->


## Checklist
- [X] Should this PR be backported?
- [X] Tests were added or are not required
- [X] Documentation was added or is not required

## Deployment Notes
<!-- Notes regarding deployment of the contained body of work. These should note any db migrations, etc. -->


Patch:
diff --git a/go/vt/vtgate/planbuilder/operator_to_query.go b/go/vt/vtgate/planbuilder/operator_to_query.go
index 864833d2c1..f14f289a3d 100644
--- a/go/vt/vtgate/planbuilder/operator_to_query.go
+++ b/go/vt/vtgate/planbuilder/operator_to_query.go
@@ -21,6 +21,8 @@ import (
 	"sort"
 	"strings"
 
+	"vitess.io/vitess/go/vt/log"
+
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/physical"
@@ -31,10 +33,16 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/abstract"
 )
 
+type queryBuilder struct {
+	ctx        *plancontext.PlanningContext
+	sel        sqlparser.SelectStatement
+	tableNames []string
+}
+
 func toSQL(ctx *plancontext.PlanningContext, op abstract.PhysicalOperator) sqlparser.SelectStatement {
 	q := &queryBuilder{ctx: ctx}
 	buildQuery(op, q)
-	q.produce()
+	q.sortTables()
 	return q.sel
 }
 
@@ -91,8 +99,20 @@ func buildQuery(op abstract.PhysicalOperator, qb *queryBuilder) {
 	}
 }
 
-func (qb *queryBuilder) produce() {
-	sort.Sort(qb)
+func (qb *queryBuilder) sortTables() {
+	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		sel, isSel := node.(*sqlparser.Select)
+		if !isSel {
+			return true, nil
+		}
+		ts := &tableSorter{
+			sel: sel,
+			tbl: qb.ctx.SemTable,
+		}
+		sort.Sort(ts)
+		return true, nil
+	}, qb.sel)
+
 }
 
 func (qb *queryBuilder) addTable(db, tableName, alias string, tableID semantics.TableSet, hints sqlparser.IndexHints) {
@@ -108,16 +128,20 @@ func (qb *queryBuilder) addTableExpr(tableName, alias string, tableID semantics.
 		qb.sel = &sqlparser.Select{}
 	}
 	sel := qb.sel.(*sqlparser.Select)
-	sel.From = append(sel.From, &sqlparser.AliasedTableExpr{
+	elems := &sqlparser.AliasedTableExpr{
 		Expr:       tblExpr,
 		Partitions: nil,
 		As:         sqlparser.NewTableIdent(alias),
 		Hints:      hints,
 		Columns:    nil,
-	})
+	}
+	err := qb.ctx.SemTable.ReplaceTableSetFor(tableID, elems)
+	if err != nil {
+		log.Warningf("error in replacing table expression in semtable: %v", err)
+	}
+	sel.From = append(sel.From, elems)
 	qb.sel = sel
 	qb.tableNames = append(qb.tableNames, tableName)
-	qb.tableIDsInFrom = append(qb.tableIDsInFrom, tableID)
 }
 
 func (qb *queryBuilder) addPredicate(expr sqlparser.Expr) {
@@ -146,7 +170,6 @@ func (qb *queryBuilder) joinInnerWith(other *queryBuilder, onCondition sqlparser
 	sel := qb.sel.(*sqlparser.Select)
 	otherSel := other.sel.(*sqlparser.Select)
 	sel.From = append(sel.From, otherSel.From...)
-	qb.tableIDsInFrom = append(qb.tableIDsInFrom, other.tableIDsInFrom...)
 	sel.SelectExprs = append(sel.SelectExprs, otherSel.SelectExprs...)
 
 	var predicate sqlparser.Expr
@@ -188,15 +211,7 @@ func (qb *queryBuilder) joinOuterWith(other *queryBuilder, onCondition sqlparser
 			On: onCondition,
 		},
 	}}
-	tableSet := semantics.EmptyTableSet()
-	for _, set := range qb.tableIDsInFrom {
-		tableSet.MergeInPlace(set)
-	}
-	for _, set := range other.tableIDsInFrom {
-		tableSet.MergeInPlace(set)
-	}
 
-	qb.tableIDsInFrom = []semantics.TableSet{tableSet}
 	sel.SelectExprs = append(sel.SelectExprs, otherSel.SelectExprs...)
 	var predicate sqlparser.Expr
 	if sel.Where != nil {
@@ -234,28 +249,33 @@ func (qb *queryBuilder) hasTable(tableName string) bool {
 	return false
 }
 
-type queryBuilder struct {
-	ctx            *plancontext.PlanningContext
-	sel            sqlparser.SelectStatement
-	tableIDsInFrom []semantics.TableSet
-	tableNames     []string
+type tableSorter struct {
+	sel *sqlparser.Select
+	tbl *semantics.SemTable
 }
 
 // Len implements the Sort interface
-func (qb *queryBuilder) Len() int {
-	return len(qb.tableIDsInFrom)
+func (ts *tableSorter) Len() int {
+	return len(ts.sel.From)
 }
 
 // Less implements the Sort interface
-func (qb *queryBuilder) Less(i, j int) bool {
-	return qb.tableIDsInFrom[i].TableOffset() < qb.tableIDsInFrom[j].TableOffset()
+func (ts *tableSorter) Less(i, j int) bool {
+	lhs := ts.sel.From[i]
+	rhs := ts.sel.From[j]
+	left, ok := lhs.(*sqlparser.AliasedTableExpr)
+	if !ok {
+		return i < j
+	}
+	right, ok := rhs.(*sqlparser.AliasedTableExpr)
+	if !ok {
+		return i < j
+	}
+
+	return ts.tbl.TableSetFor(left).TableOffset() < ts.tbl.TableSetFor(right).TableOffset()
 }
 
 // Swap implements the Sort interface
-func (qb *queryBuilder) Swap(i, j int) {
-	sel, isSel := qb.sel.(*sqlparser.Select)
-	if isSel {
-		sel.From[i], sel.From[j] = sel.From[j], sel.From[i]
-	}
-	qb.tableIDsInFrom[i], qb.tableIDsInFrom[j] = qb.tableIDsInFrom[j], qb.tableIDsInFrom[i]
+func (ts *tableSorter) Swap(i, j int) {
+	ts.sel.From[i], ts.sel.From[j] = ts.sel.From[j], ts.sel.From[i]
 }
diff --git a/go/vt/vtgate/planbuilder/single_sharded_shortcut.go b/go/vt/vtgate/planbuilder/single_sharded_shortcut.go
index efe5f928a9..5fd8ffd8fa 100644
--- a/go/vt/vtgate/planbuilder/single_sharded_shortcut.go
+++ b/go/vt/vtgate/planbuilder/single_sharded_shortcut.go
@@ -65,6 +65,10 @@ func getTableNames(semTable *semantics.SemTable) ([]string, error) {
 
 	for _, tableInfo := range semTable.Tables {
 		tblObj := tableInfo.GetVindexTable()
+		if tblObj == nil {
+			// probably a derived table
+			continue
+		}
 		var name string
 		if tableInfo.IsInfSchema() {
 			name = "tableName"
diff --git a/go/vt/vtgate/semantics/semantic_state.go b/go/vt/vtgate/semantics/semantic_state.go
index 4b33744a4c..a10d81b284 100644
--- a/go/vt/vtgate/semantics/semantic_state.go
+++ b/go/vt/vtgate/semantics/semantic_state.go
@@ -145,6 +145,26 @@ func (st *SemTable) TableSetFor(t *sqlparser.AliasedTableExpr) TableSet {
 	return TableSet{}
 }
 
+// ReplaceTableSetFor replaces the given single TabletSet with the new *sqlparser.AliasedTableExpr
+func (st *SemTable) ReplaceTableSetFor(id TableSet, t *sqlparser.AliasedTableExpr) error {
+	if id.NumberOfTables() != 1 {
+		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "BUG: tablet identifier should represent single table: %v", id)
+	}
+	tblOffset := id.TableOffset()
+	if tblOffset > len(st.Tables) {
+		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "BUG: tablet identifier greater than number of tables: %v, %d", id, len(st.Tables))
+	}
+	switch tbl := st.Tables[id.TableOffset()].(type) {
+	case *RealTable:
+		tbl.ASTNode = t
+	case *DerivedTable:
+		tbl.ASTNode = t
+	default:
+		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "BUG: replacement not expected for : %T", tbl)
+	}
+	return nil
+}
+
 // TableInfoFor returns the table info for the table set. It should contains only single table.
 func (st *SemTable) TableInfoFor(id TableSet) (TableInfo, error) {
 	offset := id.TableOffset()
@@ -351,7 +371,12 @@ func (st *SemTable) SingleUnshardedKeyspace() *vindexes.Keyspace {
 	for _, table := range st.Tables {
 		vindexTable := table.GetVindexTable()
 		if vindexTable == nil || vindexTable.Type != "" {
-			// this is not a simple table access - can't shortcut
+			_, isDT := table.getExpr().Expr.(*sqlparser.DerivedTable)
+			if isDT {
+				// derived tables are ok, as long as all real tables are from the same unsharded keyspace
+				// we check the real tables inside the derived table as well for same unsharded keyspace.
+				continue
+			}
 			return nil
 		}
 		name, ok := table.getExpr().Expr.(sqlparser.TableName)


Test patch:
diff --git a/go/vt/vtgate/planbuilder/testdata/from_cases.txt b/go/vt/vtgate/planbuilder/testdata/from_cases.txt
index 26319b68d9..85e63a7622 100644
--- a/go/vt/vtgate/planbuilder/testdata/from_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/from_cases.txt
@@ -4367,3 +4367,107 @@ Gen4 plan same as above
     "Table": "`user`"
   }
 }
+
+# single unsharded keyspace with derived table
+"select col from (select col from unsharded join unsharded_b) as u join unsharded_a ua limit 1"
+{
+  "QueryType": "SELECT",
+  "Original": "select col from (select col from unsharded join unsharded_b) as u join unsharded_a ua limit 1",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "Unsharded",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    },
+    "FieldQuery": "select col from (select col from unsharded join unsharded_b where 1 != 1) as u join unsharded_a as ua where 1 != 1",
+    "Query": "select col from (select col from unsharded join unsharded_b) as u join unsharded_a as ua limit 1",
+    "Table": "unsharded, unsharded_b, unsharded_a"
+  }
+}
+{
+  "QueryType": "SELECT",
+  "Original": "select col from (select col from unsharded join unsharded_b) as u join unsharded_a ua limit 1",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "Unsharded",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    },
+    "FieldQuery": "select col from (select col from unsharded join unsharded_b where 1 != 1) as u join unsharded_a as ua where 1 != 1",
+    "Query": "select col from (select col from unsharded join unsharded_b) as u join unsharded_a as ua limit 1",
+    "Table": "unsharded, unsharded_a, unsharded_b"
+  }
+}
+
+# query builder with derived table having join inside it
+"select u.col from (select user.col from user join user_extra) as u join user_extra ue limit 1"
+{
+  "QueryType": "SELECT",
+  "Original": "select u.col from (select user.col from user join user_extra) as u join user_extra ue limit 1",
+  "Instructions": {
+    "OperatorType": "Limit",
+    "Count": "INT64(1)",
+    "Inputs": [
+      {
+        "OperatorType": "Join",
+        "Variant": "Join",
+        "JoinColumnIndexes": "L:0",
+        "TableName": "`user`_user_extra_user_extra",
+        "Inputs": [
+          {
+            "OperatorType": "SimpleProjection",
+            "Columns": [
+              0
+            ],
+            "Inputs": [
+              {
+                "OperatorType": "Join",
+                "Variant": "Join",
+                "JoinColumnIndexes": "L:0",
+                "TableName": "`user`_user_extra",
+                "Inputs": [
+                  {
+                    "OperatorType": "Route",
+                    "Variant": "Scatter",
+                    "Keyspace": {
+                      "Name": "user",
+                      "Sharded": true
+                    },
+                    "FieldQuery": "select `user`.col from `user` where 1 != 1",
+                    "Query": "select `user`.col from `user`",
+                    "Table": "`user`"
+                  },
+                  {
+                    "OperatorType": "Route",
+                    "Variant": "Scatter",
+                    "Keyspace": {
+                      "Name": "user",
+                      "Sharded": true
+                    },
+                    "FieldQuery": "select 1 from user_extra where 1 != 1",
+                    "Query": "select 1 from user_extra",
+                    "Table": "user_extra"
+                  }
+                ]
+              }
+            ]
+          },
+          {
+            "OperatorType": "Route",
+            "Variant": "Scatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select 1 from user_extra as ue where 1 != 1",
+            "Query": "select 1 from user_extra as ue",
+            "Table": "user_extra"
+          }
+        ]
+      }
+    ]
+  }
+}
+Gen4 plan same as above