simplify merging logic (#16525)

Signed-off-by: Andres Taylor <[email protected]>

.github/workflows/vitess_tester_vtgate.yml

@@ -57,7 +57,7 @@ jobs:
           end_to_end:
             - 'go/**/*.go'
             - 'go/vt/sidecardb/**/*.sql'
-            - 'go/test/endtoend/onlineddl/vrepl_suite/**'
+            - 'go/test/endtoend/vtgate/vitess_tester/**'
             - 'test.go'
             - 'Makefile'
             - 'build.env'

go/test/endtoend/vtgate/vitess_tester/two_sharded_keyspaces/queries.test

@@ -1,26 +1,39 @@
 use customer;
-create table if not exists customer(
-                                       customer_id bigint not null,
-                                       email varbinary(128),
-                                       primary key(customer_id)
-) ENGINE=InnoDB;
-insert into customer.customer(customer_id, email) values(1, '[[email protected]](mailto:[email protected])');
-insert into customer.customer(customer_id, email) values(2, '[[email protected]](mailto:[email protected])');
-insert into customer.customer(customer_id, email) values(3, '[[email protected]](mailto:[email protected])');
-insert into customer.customer(customer_id, email) values(4, '[[email protected]](mailto:[email protected])');
-insert into customer.customer(customer_id, email) values(5, '[[email protected]](mailto:[email protected])');
+create table if not exists customer
+(
+    customer_id bigint not null,
+    email       varbinary(128),
+    primary key (customer_id)
+) ENGINE = InnoDB;
+
+insert into customer.customer(customer_id, email)
+values (1, '[[email protected]](mailto:[email protected])'),
+       (2, '[[email protected]](mailto:[email protected])'),
+       (3, '[[email protected]](mailto:[email protected])'),
+       (4, '[[email protected]](mailto:[email protected])'),
+       (5, '[[email protected]](mailto:[email protected])');
 use corder;
-create table if not exists corder(
-                                     order_id bigint not null,
-                                     customer_id bigint,
-                                     sku varbinary(128),
-                                     price bigint,
-                                     primary key(order_id)
-) ENGINE=InnoDB;
-insert into corder.corder(order_id, customer_id, sku, price) values(1, 1, 'SKU-1001', 100);
-insert into corder.corder(order_id, customer_id, sku, price) values(2, 2, 'SKU-1002', 30);
-insert into corder.corder(order_id, customer_id, sku, price) values(3, 3, 'SKU-1002', 30);
-insert into corder.corder(order_id, customer_id, sku, price) values(4, 4, 'SKU-1002', 30);
-insert into corder.corder(order_id, customer_id, sku, price) values(5, 5, 'SKU-1002', 30);
+create table if not exists corder
+(
+    order_id    bigint not null,
+    customer_id bigint,
+    sku         varbinary(128),
+    price       bigint,
+    primary key (order_id)
+) ENGINE = InnoDB;
+insert into corder.corder(order_id, customer_id, sku, price)
+values (1, 1, 'SKU-1001', 100),
+       (2, 2, 'SKU-1002', 30),
+       (3, 3, 'SKU-1002', 30),
+       (4, 4, 'SKU-1002', 30),
+       (5, 5, 'SKU-1002', 30);
+
+select co.order_id, co.customer_id, co.price
+from corder.corder co
+         left join customer.customer cu on co.customer_id = cu.customer_id
+where cu.customer_id = 1;
 
-select co.order_id, co.customer_id, co.price from corder.corder co left join customer.customer cu on co.customer_id=cu.customer_id where cu.customer_id=1;
+# This query was accidentally disallowed by https://github.com/vitessio/vitess/pull/16520
+select 1
+from customer.customer
+where customer_id in (select customer_id from corder.corder where price > 50);

go/vt/vtgate/planbuilder/operators/join_merging.go

@@ -64,7 +64,7 @@ func mergeJoinInputs(ctx *plancontext.PlanningContext, lhs, rhs Operator, joinPr
 
 	// sharded routing is complex, so we handle it in a separate method
 	case a == sharded && b == sharded:
-		return tryMergeShardedRouting(ctx, lhsRoute, rhsRoute, m, joinPredicates, false /*isSubquery*/)
+		return tryMergeShardedRouting(ctx, lhsRoute, rhsRoute, m, joinPredicates)
 
 	default:
 		return nil

go/vt/vtgate/planbuilder/operators/sharded_routing.go

@@ -23,7 +23,6 @@ import (
 	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/slice"
 	"vitess.io/vitess/go/vt/sqlparser"
-	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/engine"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
@@ -639,9 +638,10 @@ func tryMergeShardedRouting(
 	routeA, routeB *Route,
 	m merger,
 	joinPredicates []sqlparser.Expr,
-	isSubquery bool,
 ) *Route {
-	sameKeyspace := routeA.Routing.Keyspace() == routeB.Routing.Keyspace()
+	if routeA.Routing.Keyspace() != routeB.Routing.Keyspace() {
+		return nil
+	}
 	tblA := routeA.Routing.(*ShardedRouting)
 	tblB := routeB.Routing.(*ShardedRouting)
 
@@ -670,13 +670,6 @@ func tryMergeShardedRouting(
 			return nil
 		}
 
-		if !sameKeyspace {
-			if isSubquery {
-				panic(vterrors.VT12001("cross-shard correlated subquery"))
-			}
-			return nil
-		}
-
 		canMerge := canMergeOnFilters(ctx, routeA, routeB, joinPredicates)
 		if !canMerge {
 			return nil

go/vt/vtgate/planbuilder/operators/subquery_planning.go

@@ -722,7 +722,7 @@ func mergeSubqueryInputs(ctx *plancontext.PlanningContext, in, out Operator, joi
 
 	// sharded routing is complex, so we handle it in a separate method
 	case inner == sharded && outer == sharded:
-		return tryMergeShardedRouting(ctx, inRoute, outRoute, m, joinPredicates, true /*isSubquery*/)
+		return tryMergeShardedRouting(ctx, inRoute, outRoute, m, joinPredicates)
 
 	default:
 		return nil

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

@@ -296,8 +296,8 @@
   },
   {
     "comment": "Cross keyspace query with subquery",
-    "query": "select 1 from user where id in (select id from t1)",
-    "plan": "VT12001: unsupported: cross-shard correlated subquery"
+    "query": "select 1 from user where id = (select id from t1 where user.foo = t1.bar)",
+    "plan": "VT12001: unsupported: correlated subquery is only supported for EXISTS"
   },
   {
     "comment": "multi-shard union",

test/templates/cluster_vitess_tester.tpl

@@ -55,7 +55,7 @@ jobs:
           end_to_end:
             - 'go/**/*.go'
             - 'go/vt/sidecardb/**/*.sql'
-            - 'go/test/endtoend/onlineddl/vrepl_suite/**'
+            - 'go/test/endtoend/vtgate/vitess_tester/**'
             - 'test.go'
             - 'Makefile'
             - 'build.env'