v3: fix UNION merge logic

UNION mergeability was previously using the same logic as
subquery mergeability. It turns out that it doesn't always
work because subqueries can have outer query dependencies
whereas unions don't have such relationship. So, the logic
has now been separated.

Author: sougou

data/test/vtgate/select_cases.txt

@@ -565,6 +565,23 @@
   }
 }
 
+# union with the same target shard
+"select * from music where user_id = 1 union select * from user where id = 1"
+{
+  "Original": "select * from music where user_id = 1 union select * from user where id = 1",
+  "Instructions": {
+    "Opcode": "SelectEqualUnique",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "select * from music where user_id = 1 union select * from user where id = 1",
+    "FieldQuery": "select * from music where 1 != 1 union select * from user where 1 != 1",
+    "Vindex": "user_index",
+    "Values": 1
+  }
+}
+
 "select * from (select col1, col2 from unsharded where id = 1 union select col1, col2 from unsharded where id = 3) a"
 {
   "Original": "select * from (select col1, col2 from unsharded where id = 1 union select col1, col2 from unsharded where id = 3) a",

data/test/vtgate/unsupported_cases.txt

@@ -1,6 +1,6 @@
 # Unions
 "select * from user union select * from user_extra"
-"unsupported: scatter subquery"
+"unsupported: UNION on multi-shard queries"
 
 # SET
 "set a=1"
@@ -29,11 +29,11 @@
 
 # union operations in subqueries (FROM)
 "select * from (select * from user union all select * from user_extra) as t"
-"unsupported: scatter subquery"
+"unsupported: UNION on multi-shard queries"
 
 # union operations in subqueries (expressions)
 "select * from user where id in (select * from user union select * from user_extra)"
-"unsupported: scatter subquery"
+"unsupported: UNION on multi-shard queries"
 
 # subquery with join primitive (FROM)
 "select * from (select user.id from user join user_extra) as t"
@@ -265,23 +265,27 @@
 
 # multi-shard union
 "(select id from user union select id from music) union select 1 from dual"
-"unsupported: scatter subquery"
+"unsupported: UNION on multi-shard queries"
 
 # multi-shard union
 "select 1 from music union (select id from user union all select name from unsharded)"
-"unsupported: subquery keyspace different from outer query"
+"unsupported: UNION on different keyspaces"
 
 # multi-shard union
 "select 1 from music union (select id from user union select name from unsharded)"
-"unsupported: subquery keyspace different from outer query"
+"unsupported: UNION on different keyspaces"
 
 # multi-shard union
 "select id from user union all select id from music"
-"unsupported: scatter subquery"
+"unsupported: UNION on multi-shard queries"
+
+# union with different target shards
+"select 1 from music where id = 1 union select 1 from music where id = 2"
+"unsupported: UNION queries with different target shards"
 
 # Union all
 "select col1, col2 from user union all select col1, col2 from user_extra"
-"unsupported: scatter subquery"
+"unsupported: UNION on multi-shard queries"
 
 "(select user.id, user.name from user join user_extra where user_extra.extra = 'asdf') union select 'b','c' from user"
 "unsupported construct: SELECT of UNION is non-trivial"

go/vt/vtgate/planbuilder/expr.go

@@ -93,7 +93,7 @@ func findRoute(expr sqlparser.Expr, bldr builder) (rb *route, err error) {
 		return nil, err
 	}
 	for _, subroute := range subroutes {
-		err = routesCanMerge(highestRoute, subroute)
+		err = subqueryCanMerge(highestRoute, subroute)
 		if err != nil {
 			return nil, err
 		}
@@ -104,10 +104,10 @@ func findRoute(expr sqlparser.Expr, bldr builder) (rb *route, err error) {
 	return highestRoute, nil
 }
 
-// routesCanMerge returns nil if the inner subquery
+// subqueryCanMerge returns nil if the inner subquery
 // can be merged with the specified outer route. If it
 // cannot, then it returns an appropriate error.
-func routesCanMerge(outer, inner *route) error {
+func subqueryCanMerge(outer, inner *route) error {
 	if outer.ERoute.Keyspace.Name != inner.ERoute.Keyspace.Name {
 		return errors.New("unsupported: subquery keyspace different from outer query")
 	}

go/vt/vtgate/planbuilder/union.go

@@ -95,3 +95,25 @@ func unionRouteMerge(union *sqlparser.Union, left, right builder, vschema VSchem
 	rroute.Redirect = rtb
 	return rtb, nil
 }
+
+// routesCanMerge returns nil if the left and right route
+// can be merged. Otherwise, it returns an appropriate error.
+func routesCanMerge(left, right *route) error {
+	if left.ERoute.Opcode == engine.SelectNext || right.ERoute.Opcode == engine.SelectNext {
+		return errors.New("unsupported: UNION on sequence tables")
+	}
+	if left.ERoute.Keyspace.Name != right.ERoute.Keyspace.Name {
+		return errors.New("unsupported: UNION on different keyspaces")
+	}
+	if left.ERoute.Opcode == engine.SelectUnsharded {
+		// right will also be unsharded. So, we're good.
+		return nil
+	}
+	if left.ERoute.Opcode != engine.SelectEqualUnique || right.ERoute.Opcode != engine.SelectEqualUnique {
+		return errors.New("unsupported: UNION on multi-shard queries")
+	}
+	if !valEqual(left.ERoute.Values, right.ERoute.Values) {
+		return errors.New("unsupported: UNION queries with different target shards")
+	}
+	return nil
+}