Fix vtexplain not handling UNION queries with weight_string res…

…ults correctly. (#16129)

Signed-off-by: Arthur Schreiber <[email protected]>

go/vt/vtexplain/testdata/multi-output/selectsharded-output.txt

@@ -91,18 +91,18 @@ select name, count(*) from user group by name /* scatter aggregate */
 ----------------------------------------------------------------------
 select 1, "hello", 3.14, null from user limit 10 /* select constant sql values */
 
-1 ks_sharded/-40: select 1, 'hello', 3.14, null from `user` limit 10 /* select constant sql values */
-1 ks_sharded/40-80: select 1, 'hello', 3.14, null from `user` limit 10 /* select constant sql values */
-1 ks_sharded/80-c0: select 1, 'hello', 3.14, null from `user` limit 10 /* select constant sql values */
-1 ks_sharded/c0-: select 1, 'hello', 3.14, null from `user` limit 10 /* select constant sql values */
+1 ks_sharded/-40: select 1, 'hello', 3.14, null from `user` limit 10 /* INT64 */ /* select constant sql values */
+1 ks_sharded/40-80: select 1, 'hello', 3.14, null from `user` limit 10 /* INT64 */ /* select constant sql values */
+1 ks_sharded/80-c0: select 1, 'hello', 3.14, null from `user` limit 10 /* INT64 */ /* select constant sql values */
+1 ks_sharded/c0-: select 1, 'hello', 3.14, null from `user` limit 10 /* INT64 */ /* select constant sql values */
 
 ----------------------------------------------------------------------
 select * from (select id from user) s /* scatter paren select */
 
-1 ks_sharded/-40: select s.id from (select id from `user`) as s limit 10001 /* scatter paren select */
-1 ks_sharded/40-80: select s.id from (select id from `user`) as s limit 10001 /* scatter paren select */
-1 ks_sharded/80-c0: select s.id from (select id from `user`) as s limit 10001 /* scatter paren select */
-1 ks_sharded/c0-: select s.id from (select id from `user`) as s limit 10001 /* scatter paren select */
+1 ks_sharded/-40: select id from (select id from `user`) as s limit 10001 /* scatter paren select */
+1 ks_sharded/40-80: select id from (select id from `user`) as s limit 10001 /* scatter paren select */
+1 ks_sharded/80-c0: select id from (select id from `user`) as s limit 10001 /* scatter paren select */
+1 ks_sharded/c0-: select id from (select id from `user`) as s limit 10001 /* scatter paren select */
 
 ----------------------------------------------------------------------
 select name from user where id = (select id from t1) /* non-correlated subquery as value */
@@ -114,21 +114,21 @@ select name from user where id = (select id from t1) /* non-correlated subquery
 select name from user where id in (select id from t1) /* non-correlated subquery in IN clause */
 
 1 ks_unsharded/-: select id from t1 limit 10001 /* non-correlated subquery in IN clause */
-2 ks_sharded/-40: select `name` from `user` where 1 = 1 and id in (1) limit 10001 /* non-correlated subquery in IN clause */
+2 ks_sharded/-40: select `name` from `user` where 1 and id in (1) limit 10001 /* non-correlated subquery in IN clause */
 
 ----------------------------------------------------------------------
 select name from user where id not in (select id from t1) /* non-correlated subquery in NOT IN clause */
 
 1 ks_unsharded/-: select id from t1 limit 10001 /* non-correlated subquery in NOT IN clause */
-2 ks_sharded/-40: select `name` from `user` where 1 = 0 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
-2 ks_sharded/40-80: select `name` from `user` where 1 = 0 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
-2 ks_sharded/80-c0: select `name` from `user` where 1 = 0 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
-2 ks_sharded/c0-: select `name` from `user` where 1 = 0 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
+2 ks_sharded/-40: select `name` from `user` where not 1 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
+2 ks_sharded/40-80: select `name` from `user` where not 1 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
+2 ks_sharded/80-c0: select `name` from `user` where not 1 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
+2 ks_sharded/c0-: select `name` from `user` where not 1 or id not in (1) limit 10001 /* non-correlated subquery in NOT IN clause */
 
 ----------------------------------------------------------------------
 select name from user where exists (select id from t1) /* non-correlated subquery as EXISTS */
 
-1 ks_unsharded/-: select 1 from t1 limit 1 /* non-correlated subquery as EXISTS */
+1 ks_unsharded/-: select 1 from t1 limit 10001 /* non-correlated subquery as EXISTS */
 2 ks_sharded/-40: select `name` from `user` where 1 limit 10001 /* non-correlated subquery as EXISTS */
 2 ks_sharded/40-80: select `name` from `user` where 1 limit 10001 /* non-correlated subquery as EXISTS */
 2 ks_sharded/80-c0: select `name` from `user` where 1 limit 10001 /* non-correlated subquery as EXISTS */
@@ -137,10 +137,10 @@ select name from user where exists (select id from t1) /* non-correlated subquer
 ----------------------------------------------------------------------
 select * from name_info order by info /* select * and order by varchar column */
 
-1 ks_sharded/-40: select `name`, info from name_info order by info asc limit 10001 /* select * and order by varchar column */
-1 ks_sharded/40-80: select `name`, info from name_info order by info asc limit 10001 /* select * and order by varchar column */
-1 ks_sharded/80-c0: select `name`, info from name_info order by info asc limit 10001 /* select * and order by varchar column */
-1 ks_sharded/c0-: select `name`, info from name_info order by info asc limit 10001 /* select * and order by varchar column */
+1 ks_sharded/-40: select `name`, info, weight_string(info) from name_info order by name_info.info asc limit 10001 /* select * and order by varchar column */
+1 ks_sharded/40-80: select `name`, info, weight_string(info) from name_info order by name_info.info asc limit 10001 /* select * and order by varchar column */
+1 ks_sharded/80-c0: select `name`, info, weight_string(info) from name_info order by name_info.info asc limit 10001 /* select * and order by varchar column */
+1 ks_sharded/c0-: select `name`, info, weight_string(info) from name_info order by name_info.info asc limit 10001 /* select * and order by varchar column */
 
 ----------------------------------------------------------------------
 select distinct(name) from user where id = 1 /* select distinct */
@@ -207,3 +207,9 @@ SELECT id FROM orders WHERE id IN (1, "1", 1)
 2 ks_sharded/40-80: select id from orders where id in (1, '1', 1) limit 10001
 
 ----------------------------------------------------------------------
+(SELECT user.id, user.name FROM user WHERE user.id = 1) UNION (SELECT user.id, user.name FROM user WHERE user.id = 2)
+
+2 ks_sharded/-40: select dt.c0 as id, dt.c1 as `name`, weight_string(dt.c0), weight_string(dt.c1) from (select distinct `user`.id, `user`.`name` from `user` where `user`.id = 2) as dt(c0, c1) limit 10001
+2 ks_sharded/-40: select dt.c0 as id, dt.c1 as `name`, weight_string(dt.c0), weight_string(dt.c1) from (select distinct `user`.id, `user`.`name` from `user` where `user`.id = 1) as dt(c0, c1) limit 10001
+
+----------------------------------------------------------------------

go/vt/vtexplain/testdata/multi-output/updatesharded-output.txt

@@ -128,10 +128,10 @@ update user set nickname='alice' where id in (1,4)
 
 1 ks_sharded/-40: begin
 1 ks_sharded/-40: savepoint x1
-1 ks_sharded/-40: update `user` set nickname = 'alice' where id in (1, 4) limit 10001
+1 ks_sharded/-40: update `user` set nickname = 'alice' where id in (1) limit 10001
 1 ks_sharded/c0-: begin
 1 ks_sharded/c0-: savepoint x1
-1 ks_sharded/c0-: update `user` set nickname = 'alice' where id in (1, 4) limit 10001
+1 ks_sharded/c0-: update `user` set nickname = 'alice' where id in (4) limit 10001
 
 ----------------------------------------------------------------------
 commit

go/vt/vtexplain/testdata/selectsharded-queries.sql

@@ -38,4 +38,6 @@ select id from user where not id in (select col from music where music.user_id =
 
 SELECT user.id, user.name, name_info.info FROM user INNER JOIN music ON (user.id = music.user_id) LEFT OUTER JOIN name_info ON (user.name = name_info.name);
 
-SELECT id FROM orders WHERE id IN (1, "1", 1)
+SELECT id FROM orders WHERE id IN (1, "1", 1);
+
+(SELECT user.id, user.name FROM user WHERE user.id = 1) UNION (SELECT user.id, user.name FROM user WHERE user.id = 2);

go/vt/vtexplain/vtexplain_test.go

@@ -35,6 +35,7 @@ import (
 	"vitess.io/vitess/go/test/utils"
 
 	"vitess.io/vitess/go/vt/key"
+	querypb "vitess.io/vitess/go/vt/proto/query"
 	"vitess.io/vitess/go/vt/proto/topodata"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv/tabletenvtest"
@@ -151,6 +152,28 @@ func TestExplain(t *testing.T) {
 	}
 	tests := []test{
 		{"unsharded", defaultTestOpts()},
+		{"selectsharded", defaultTestOpts()},
+		{"insertsharded", defaultTestOpts()},
+		{"updatesharded", defaultTestOpts()},
+		{"deletesharded", defaultTestOpts()},
+		{"comments", defaultTestOpts()},
+		{"options", &Options{
+			ReplicationMode: "STATEMENT",
+			NumShards:       4,
+			Normalize:       false,
+		}},
+		{"target", &Options{
+			ReplicationMode: "ROW",
+			NumShards:       4,
+			Normalize:       false,
+			Target:          "ks_sharded/40-80",
+		}},
+		{"gen4", &Options{
+			ReplicationMode: "ROW",
+			NumShards:       4,
+			Normalize:       true,
+			PlannerVersion:  querypb.ExecuteOptions_Gen4,
+		}},
 	}
 
 	for _, tst := range tests {

go/vt/vtexplain/vtexplain_vttablet.go

@@ -864,9 +864,15 @@ func inferColTypeFromExpr(node sqlparser.Expr, tableColumnMap map[sqlparser.Iden
 			colTypes = append(colTypes, colType)
 		}
 	case sqlparser.Callable:
-		// As a shortcut, functions are integral types
-		colNames = append(colNames, sqlparser.String(node))
-		colTypes = append(colTypes, querypb.Type_INT32)
+		switch node := node.(type) {
+		case *sqlparser.WeightStringFuncExpr:
+			colNames = append(colNames, sqlparser.String(node))
+			colTypes = append(colTypes, querypb.Type_BINARY)
+		default:
+			// As a shortcut, functions are integral types
+			colNames = append(colNames, sqlparser.String(node))
+			colTypes = append(colTypes, querypb.Type_INT32)
+		}
 	case *sqlparser.Literal:
 		colNames = append(colNames, sqlparser.String(node))
 		switch node.Type {