gen4: support for union in subquery

go/vt/sqlparser/ast_format.go

@@ -1728,7 +1728,9 @@ func (node *RenameTable) Format(buf *TrackedBuffer) {
 func (node *ExtractedSubquery) Format(buf *TrackedBuffer) {
 	switch original := node.Original.(type) {
 	case *ExistsExpr:
-		buf.astPrintf(node, "%v", NewArgument(node.ArgName))
+		buf.formatter(NewArgument(node.ArgName))
+	case *Subquery:
+		buf.formatter(NewArgument(node.ArgName))
 	case *ComparisonExpr:
 		// other_side = :__sq
 		cmp := &ComparisonExpr{
@@ -1749,8 +1751,6 @@ func (node *ExtractedSubquery) Format(buf *TrackedBuffer) {
 			hasValue := &ComparisonExpr{Left: NewArgument(node.HasValuesArg), Right: NewIntLiteral("0"), Operator: EqualOp}
 			expr = &OrExpr{hasValue, cmp}
 		}
-		buf.astPrintf(node, "%v", expr)
-	case *Subquery:
-		buf.astPrintf(node, "%v", NewArgument(node.ArgName))
+		buf.formatter(expr)
 	}
 }

go/vt/vtgate/planbuilder/abstract/concatenate.go

@@ -50,7 +50,7 @@ func (c *Concatenate) PushPredicate(sqlparser.Expr, *semantics.SemTable) error {
 
 // UnsolvedPredicates implements the Operator interface
 func (c *Concatenate) UnsolvedPredicates(*semantics.SemTable) []sqlparser.Expr {
-	panic("implement me")
+	return nil
 }
 
 // CheckValid implements the Operator interface

go/vt/vtgate/planbuilder/abstract/operator.go

@@ -190,17 +190,13 @@ func createOperatorFromSelect(sel *sqlparser.Select, semTable *semantics.SemTabl
 	if len(semTable.SubqueryMap[sel]) > 0 {
 		resultantOp = &SubQuery{}
 		for _, sq := range semTable.SubqueryMap[sel] {
-			subquerySelectStatement, isSel := sq.Subquery.Select.(*sqlparser.Select)
-			if !isSel {
-				return nil, semantics.Gen4NotSupportedF("UNION in subquery")
-			}
-			opInner, err := createOperatorFromSelect(subquerySelectStatement, semTable)
+			opInner, err := CreateOperatorFromAST(sq.Subquery.Select, semTable)
 			if err != nil {
 				return nil, err
 			}
 			resultantOp.Inner = append(resultantOp.Inner, &SubQueryInner{
-				Inner:             opInner,
 				ExtractedSubquery: sq,
+				Inner:             opInner,
 			})
 		}
 	}

go/vt/vtgate/planbuilder/querytree_transformers.go

@@ -98,15 +98,43 @@ func transformSubqueryTree(ctx *planningContext, n *subqueryTree) (logicalPlan,
 		return nil, err
 	}
 
-	plan := newPulloutSubquery(engine.PulloutOpcode(n.extracted.OpCode), n.extracted.ArgName, n.extracted.HasValuesArg, innerPlan)
+	argName := n.extracted.ArgName
+	hasValuesArg := n.extracted.HasValuesArg
 	outerPlan, err := transformToLogicalPlan(ctx, n.outer)
+
+	merged := mergeSubQueryPlan(ctx, innerPlan, outerPlan, n)
+	if merged != nil {
+		return merged, nil
+	}
+	plan := newPulloutSubquery(engine.PulloutOpcode(n.extracted.OpCode), argName, hasValuesArg, innerPlan)
 	if err != nil {
 		return nil, err
 	}
 	plan.underlying = outerPlan
 	return plan, err
 }
 
+func mergeSubQueryPlan(ctx *planningContext, inner, outer logicalPlan, n *subqueryTree) logicalPlan {
+	iroute, ok := inner.(*route)
+	if !ok {
+		return nil
+	}
+	oroute, ok := outer.(*route)
+	if !ok {
+		return nil
+	}
+
+	if canMergeSubqueryPlans(ctx, iroute, oroute) {
+		// n.extracted is an expression that lives in oroute.Select.
+		// 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 nil
+}
+
 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
@@ -267,7 +295,7 @@ func mergeUnionLogicalPlans(ctx *planningContext, left logicalPlan, right logica
 		return nil
 	}
 
-	if canMergePlans(ctx, lroute, rroute) {
+	if canMergeUnionPlans(ctx, lroute, rroute) {
 		lroute.Select = &sqlparser.Union{Left: lroute.Select, Distinct: false, Right: rroute.Select}
 		return lroute
 	}
@@ -522,7 +550,7 @@ func (sqr *subQReplacer) replacer(cursor *sqlparser.Cursor) bool {
 	return true
 }
 
-func replaceSubQuery(ctx *planningContext, sel *sqlparser.Select) {
+func replaceSubQuery(ctx *planningContext, sel sqlparser.SelectStatement) {
 	extractedSubqueries := ctx.semTable.GetSubqueryNeedingRewrite()
 	if len(extractedSubqueries) == 0 {
 		return

go/vt/vtgate/planbuilder/rewrite_test.go

@@ -41,7 +41,7 @@ func TestSubqueryRewrite(t *testing.T) {
 		output: "select 1 from t1 where :__sq_has_values1",
 	}, {
 		input:  "select id from t1 where id in (select 1)",
-		output: "select id from t1 where (:__sq_has_values1 = 1 and id in ::__sq1)",
+		output: "select id from t1 where :__sq_has_values1 = 1 and id in ::__sq1",
 	}, {
 		input:  "select id from t1 where id not in (select 1)",
 		output: "select id from t1 where :__sq_has_values1 = 0 or id not in ::__sq1",
@@ -65,7 +65,7 @@ func TestSubqueryRewrite(t *testing.T) {
 		output: "select id from t1 where not :__sq_has_values1 and :__sq_has_values2",
 	}, {
 		input:  "select (select 1), (select 2) from t1 join t2 on t1.id = (select 1) where t1.id in (select 1)",
-		output: "select :__sq2, :__sq3 from t1 join t2 on t1.id = :__sq1 where (:__sq_has_values4 = 1 and t1.id in ::__sq4)",
+		output: "select :__sq2, :__sq3 from t1 join t2 on t1.id = :__sq1 where :__sq_has_values4 = 1 and t1.id in ::__sq4",
 	}}
 	for _, tcase := range tcases {
 		t.Run(tcase.input, func(t *testing.T) {
@@ -123,7 +123,7 @@ func TestHavingRewrite(t *testing.T) {
 		output: "select count(*) as k from t1 where (x = 1 or y = 2) and a = 1 having count(*) = 1",
 	}, {
 		input:  "select 1 from t1 where x in (select 1 from t2 having a = 1)",
-		output: "select 1 from t1 where (:__sq_has_values1 = 1 and x in ::__sq1)",
+		output: "select 1 from t1 where :__sq_has_values1 = 1 and x in ::__sq1",
 		sqs:    map[string]string{"__sq1": "select 1 from t2 where a = 1"},
 	}, {input: "select 1 from t1 group by a having a = 1 and count(*) > 1",
 		output: "select 1 from t1 where a = 1 group by a having count(*) > 1",

go/vt/vtgate/planbuilder/route_planning.go

@@ -889,7 +889,7 @@ func canMergeOnFilters(ctx *planningContext, a, b *routeTree, joinPredicates []s
 
 type mergeFunc func(a, b *routeTree) (*routeTree, error)
 
-func canMergePlans(ctx *planningContext, a, b *route) bool {
+func canMergeUnionPlans(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 {
 		return false
@@ -919,6 +919,32 @@ func canMergePlans(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 {
+		return false
+	}
+	switch a.eroute.Opcode {
+	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
+	case engine.SelectEqualUnique:
+		// Check if they target the same shard.
+		if b.eroute.Opcode == engine.SelectEqualUnique &&
+			a.eroute.Vindex == b.eroute.Vindex &&
+			a.condition != nil &&
+			b.condition != nil &&
+			gen4ValuesEqual(ctx, []sqlparser.Expr{a.condition}, []sqlparser.Expr{b.condition}) {
+			return true
+		}
+	}
+	return false
+}
 
 func tryMerge(ctx *planningContext, a, b queryTree, joinPredicates []sqlparser.Expr, merger mergeFunc) (queryTree, error) {
 	aRoute, bRoute := queryTreesToRoutes(a.clone(), b.clone())

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

@@ -264,6 +264,7 @@ Gen4 plan same as above
     "Query": "create view view_a as select id, `name` from unsharded where id in (select id from unsharded where id = 1 union select id from unsharded where id = 3)"
   }
 }
+Gen4 plan same as above
 
 # create view with subquery in unsharded keyspace with UNION clause
 "create view view_a as (select id from unsharded) union (select id from unsharded_auto) order by id limit 5"

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

@@ -1738,46 +1738,7 @@ Gen4 plan same as above
     ]
   }
 }
-{
-  "QueryType": "SELECT",
-  "Original": "select id from user where id in (select col from user)",
-  "Instructions": {
-    "OperatorType": "Subquery",
-    "Variant": "PulloutIn",
-    "PulloutVars": [
-      "__sq_has_values1",
-      "__sq1"
-    ],
-    "Inputs": [
-      {
-        "OperatorType": "Route",
-        "Variant": "SelectScatter",
-        "Keyspace": {
-          "Name": "user",
-          "Sharded": true
-        },
-        "FieldQuery": "select col from `user` where 1 != 1",
-        "Query": "select col from `user`",
-        "Table": "`user`"
-      },
-      {
-        "OperatorType": "Route",
-        "Variant": "SelectIN",
-        "Keyspace": {
-          "Name": "user",
-          "Sharded": true
-        },
-        "FieldQuery": "select id from `user` where 1 != 1",
-        "Query": "select id from `user` where (:__sq_has_values1 = 1 and id in ::__vals)",
-        "Table": "`user`",
-        "Values": [
-          "::__sq1"
-        ],
-        "Vindex": "user_index"
-      }
-    ]
-  }
-}
+Gen4 plan same as above
 
 # cross-shard subquery in NOT IN clause.
 "select id from user where id not in (select col from user)"
@@ -2038,7 +1999,7 @@ Gen4 plan same as above
               "Sharded": true
             },
             "FieldQuery": "select id2 from `user` where 1 != 1",
-            "Query": "select id2 from `user` where (:__sq_has_values2 = 1 and id2 in ::__sq2)",
+            "Query": "select id2 from `user` where :__sq_has_values2 = 1 and id2 in ::__sq2",
             "Table": "`user`"
           }
         ]
@@ -2473,7 +2434,7 @@ Gen4 plan same as above
               "Sharded": true
             },
             "FieldQuery": "select id from `user` where 1 != 1",
-            "Query": "select id from `user` where :__sq_has_values1 = 0 or id not in ::__sq1 and (:__sq_has_values2 = 1 and id in ::__vals)",
+            "Query": "select id from `user` where :__sq_has_values1 = 0 or id not in ::__sq1 and :__sq_has_values2 = 1 and id in ::__vals",
             "Table": "`user`",
             "Values": [
               "::__sq2"
@@ -2611,46 +2572,7 @@ Gen4 plan same as above
     ]
   }
 }
-{
-  "QueryType": "SELECT",
-  "Original": "select id from user where id in (select col from unsharded where col = id)",
-  "Instructions": {
-    "OperatorType": "Subquery",
-    "Variant": "PulloutIn",
-    "PulloutVars": [
-      "__sq_has_values1",
-      "__sq1"
-    ],
-    "Inputs": [
-      {
-        "OperatorType": "Route",
-        "Variant": "SelectUnsharded",
-        "Keyspace": {
-          "Name": "main",
-          "Sharded": false
-        },
-        "FieldQuery": "select col from unsharded where 1 != 1",
-        "Query": "select col from unsharded where col = id",
-        "Table": "unsharded"
-      },
-      {
-        "OperatorType": "Route",
-        "Variant": "SelectIN",
-        "Keyspace": {
-          "Name": "user",
-          "Sharded": true
-        },
-        "FieldQuery": "select id from `user` where 1 != 1",
-        "Query": "select id from `user` where (:__sq_has_values1 = 1 and id in ::__vals)",
-        "Table": "`user`",
-        "Values": [
-          "::__sq1"
-        ],
-        "Vindex": "user_index"
-      }
-    ]
-  }
-}
+Gen4 plan same as above
 
 # correlated subquery with different keyspace tables involved
 "select id from user where id in (select col from unsharded where col = user.id)"
@@ -2945,7 +2867,7 @@ Gen4 plan same as above
           "Sharded": true
         },
         "FieldQuery": "select id from `user` where 1 != 1",
-        "Query": "select id from `user` where id = 5 and id not in (select user_extra.col from user_extra where user_extra.user_id = 5) and (:__sq_has_values2 = 1 and id in ::__sq2)",
+        "Query": "select id from `user` where id = 5 and id not in (select user_extra.col from user_extra where user_extra.user_id = 5) and :__sq_has_values2 = 1 and id in ::__sq2",
         "Table": "`user`",
         "Values": [
           5
@@ -3393,7 +3315,7 @@ Gen4 plan same as above
             },
             "FieldQuery": "select `user`.id, `user`.col, weight_string(`user`.id), weight_string(`user`.col) from `user` where 1 != 1",
             "OrderBy": "(0|2) ASC, (1|3) ASC",
-            "Query": "select `user`.id, `user`.col, weight_string(`user`.id), weight_string(`user`.col) from `user` where (:__sq_has_values1 = 1 and `user`.col in ::__sq1) order by `user`.id asc, `user`.col asc",
+            "Query": "select `user`.id, `user`.col, weight_string(`user`.id), weight_string(`user`.col) from `user` where :__sq_has_values1 = 1 and `user`.col in ::__sq1 order by `user`.id asc, `user`.col asc",
             "Table": "`user`"
           }
         ]