Gen4: Improve merging of DBA queries

go/vt/vtgate/planbuilder/horizon_planning.go

@@ -316,6 +316,18 @@ func pushProjection(expr *sqlparser.AliasedExpr, plan logicalPlan, semTable *sem
 		}
 		node.cols = append(node.cols, column)
 		return len(node.cols) - 1, true, nil
+	case *concatenateGen4:
+		if hasAggregation {
+			return 0, false, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: aggregation on unions")
+		}
+		offset, added, err := pushProjection(expr, node.sources[0], semTable, inner, reuseCol, hasAggregation)
+		if err != nil {
+			return 0, false, err
+		}
+		if added {
+			return 0, false, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "pushing projection %v on concatenate should reference an existing column", sqlparser.String(expr))
+		}
+		return offset, false, nil
 	default:
 		return 0, false, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "[BUG] push projection does not yet support: %T", node)
 	}

go/vt/vtgate/planbuilder/querytree_transformers.go

@@ -143,12 +143,23 @@ func mergeSubQueryPlan(ctx *planningContext, inner, outer logicalPlan, n *subque
 		// Instead of looking for it in the AST, we have a copy in the subquery tree that we can update
 		n.extracted.NeedsRewrite = true
 		replaceSubQuery(ctx, oroute.Select)
-
-		return oroute
+		return mergeSystemTableInformation(oroute, iroute)
 	}
 	return nil
 }
 
+// mergeSystemTableInformation copies over information from the second route to the first and appends to it
+func mergeSystemTableInformation(a *route, b *route) logicalPlan {
+	// safe to append system table schema and system table names, since either the routing will match or either side would be throwing an error
+	// during run-time which we want to preserve. For example outer side has User in sys table schema and inner side has User and Main in sys table schema
+	// Inner might end up throwing an error at runtime, but if it doesn't then it is safe to merge.
+	a.eroute.SysTableTableSchema = append(a.eroute.SysTableTableSchema, b.eroute.SysTableTableSchema...)
+	for k, v := range b.eroute.SysTableTableName {
+		a.eroute.SysTableTableName[k] = v
+	}
+	return a
+}
+
 func transformDerivedPlan(ctx *planningContext, n *derivedTree) (logicalPlan, error) {
 	// transforming the inner part of the derived table into a logical plan
 	// so that we can do horizon planning on the inner. If the logical plan
@@ -311,7 +322,7 @@ func mergeUnionLogicalPlans(ctx *planningContext, left logicalPlan, right logica
 
 	if canMergeUnionPlans(ctx, lroute, rroute) {
 		lroute.Select = &sqlparser.Union{Left: lroute.Select, Distinct: false, Right: rroute.Select}
-		return lroute
+		return mergeSystemTableInformation(lroute, rroute)
 	}
 	return nil
 }

go/vt/vtgate/planbuilder/route_planning.go

@@ -971,11 +971,7 @@ func canMergeUnionPlans(ctx *planningContext, a, b *route) bool {
 	case engine.SelectUnsharded, engine.SelectReference:
 		return a.eroute.Opcode == b.eroute.Opcode
 	case engine.SelectDBA:
-		return b.eroute.Opcode == engine.SelectDBA &&
-			len(a.eroute.SysTableTableSchema) == 0 &&
-			len(a.eroute.SysTableTableName) == 0 &&
-			len(b.eroute.SysTableTableSchema) == 0 &&
-			len(b.eroute.SysTableTableName) == 0
+		return canSelectDBAMerge(a, b)
 	case engine.SelectEqualUnique:
 		// Check if they target the same shard.
 		if b.eroute.Opcode == engine.SelectEqualUnique &&
@@ -992,6 +988,7 @@ func canMergeUnionPlans(ctx *planningContext, a, b *route) bool {
 	}
 	return false
 }
+
 func canMergeSubqueryPlans(ctx *planningContext, a, b *route) bool {
 	// this method should be close to tryMerge below. it does the same thing, but on logicalPlans instead of queryTrees
 	if a.eroute.Keyspace.Name != b.eroute.Keyspace.Name {
@@ -1001,11 +998,7 @@ func canMergeSubqueryPlans(ctx *planningContext, a, b *route) bool {
 	case engine.SelectUnsharded, engine.SelectReference:
 		return a.eroute.Opcode == b.eroute.Opcode
 	case engine.SelectDBA:
-		return b.eroute.Opcode == engine.SelectDBA &&
-			len(a.eroute.SysTableTableSchema) == 0 &&
-			len(a.eroute.SysTableTableName) == 0 &&
-			len(b.eroute.SysTableTableSchema) == 0 &&
-			len(b.eroute.SysTableTableName) == 0
+		return canSelectDBAMerge(a, b)
 	case engine.SelectEqualUnique:
 		// Check if they target the same shard.
 		if b.eroute.Opcode == engine.SelectEqualUnique &&
@@ -1019,6 +1012,37 @@ func canMergeSubqueryPlans(ctx *planningContext, a, b *route) bool {
 	return false
 }
 
+func canSelectDBAMerge(a, b *route) bool {
+	if a.eroute.Opcode != engine.SelectDBA {
+		return false
+	}
+	if b.eroute.Opcode != engine.SelectDBA {
+		return false
+	}
+
+	// safe to merge when any 1 table name or schema matches, since either the routing will match or either side would be throwing an error
+	// during run-time which we want to preserve. For example outer side has User in sys table schema and inner side has User and Main in sys table schema
+	// Inner might end up throwing an error at runtime, but if it doesn't then it is safe to merge.
+	for _, aExpr := range a.eroute.SysTableTableSchema {
+		for _, bExpr := range b.eroute.SysTableTableSchema {
+			if aExpr.String() == bExpr.String() {
+				return true
+			}
+		}
+	}
+	for _, aExpr := range a.eroute.SysTableTableName {
+		for _, bExpr := range b.eroute.SysTableTableName {
+			if aExpr.String() == bExpr.String() {
+				return true
+			}
+		}
+	}
+
+	// if either/both of the side does not have any routing information, then they can be merged.
+	return (len(a.eroute.SysTableTableSchema) == 0 && len(a.eroute.SysTableTableName) == 0) ||
+		(len(b.eroute.SysTableTableSchema) == 0 && len(b.eroute.SysTableTableName) == 0)
+}
+
 func tryMerge(ctx *planningContext, a, b queryTree, joinPredicates []sqlparser.Expr, merger mergeFunc) (queryTree, error) {
 	aRoute, bRoute := queryTreesToRoutes(a.clone(), b.clone())
 	if aRoute == nil || bRoute == nil {

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

@@ -2730,3 +2730,46 @@ Gen4 plan same as above
 "select 1 from user u where u.col = 6 or exists (select 1 from user_extra ue where ue.col = u.col and u.col = ue.col2)"
 "unsupported: cross-shard correlated subquery"
 Gen4 error: exists sub-queries are only supported with AND clause
+
+# union as a derived table
+"select found from (select id as found from user union all (select id from unsharded)) as t"
+{
+  "QueryType": "SELECT",
+  "Original": "select found from (select id as found from user union all (select id from unsharded)) as t",
+  "Instructions": {
+    "OperatorType": "SimpleProjection",
+    "Columns": [
+      0
+    ],
+    "Inputs": [
+      {
+        "OperatorType": "Concatenate",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select id as found from `user` where 1 != 1",
+            "Query": "select id as found from `user`",
+            "Table": "`user`"
+          },
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectUnsharded",
+            "Keyspace": {
+              "Name": "main",
+              "Sharded": false
+            },
+            "FieldQuery": "select id from unsharded where 1 != 1",
+            "Query": "select id from unsharded",
+            "Table": "unsharded"
+          }
+        ]
+      }
+    ]
+  }
+}
+Gen4 plan same as above