diff --git a/go/vt/vtgate/planbuilder/operators/horizon.go b/go/vt/vtgate/planbuilder/operators/horizon.go index e576a1c4cba..f78448dc257 100644 --- a/go/vt/vtgate/planbuilder/operators/horizon.go +++ b/go/vt/vtgate/planbuilder/operators/horizon.go @@ -119,7 +119,18 @@ func (h *Horizon) AddColumn(ctx *plancontext.PlanningContext, reuse bool, _ bool } func (h *Horizon) AddWSColumn(ctx *plancontext.PlanningContext, offset int, underRoute bool) int { - panic(errNoNewColumns) + cols := h.GetColumns(ctx) + if offset >= len(cols) { + panic(errNoNewColumns) + } + + sel, ok := h.Query.(*sqlparser.Select) + if !ok { + panic(errNoNewColumns) + } + wsOffset := len(cols) + sel.AddSelectExpr(aeWrap(weightStringFor(cols[offset].Expr))) + return wsOffset } var errNoNewColumns = vterrors.VT13001("can't add new columns to Horizon") diff --git a/go/vt/vtgate/planbuilder/operators/route.go b/go/vt/vtgate/planbuilder/operators/route.go index 43cb22842ad..38149b5dac5 100644 --- a/go/vt/vtgate/planbuilder/operators/route.go +++ b/go/vt/vtgate/planbuilder/operators/route.go @@ -692,6 +692,16 @@ func addWSColumnToInput(ctx *plancontext.PlanningContext, source Operator, offse return true, op.AddWSColumn(ctx, offset, true) case *Aggregator: return true, op.AddWSColumn(ctx, offset, true) + case *Union: + cloned := slice.Map(op.Sources, func(src Operator) Operator { + return Clone(src) + }) + wsOffset, err := op.addWeightStringToOffset(ctx, offset) + if err != nil { + op.Sources = cloned + return false, -1 + } + return true, wsOffset } return false, -1 } diff --git a/go/vt/vtgate/planbuilder/operators/union.go b/go/vt/vtgate/planbuilder/operators/union.go index 0c7ecc97bc3..ae72244ca07 100644 --- a/go/vt/vtgate/planbuilder/operators/union.go +++ b/go/vt/vtgate/planbuilder/operators/union.go @@ -160,7 +160,11 @@ func (u *Union) GetSelectFor(source int) *sqlparser.Select { } func (u *Union) AddWSColumn(ctx *plancontext.PlanningContext, offset int, underRoute bool) int { - return u.addWeightStringToOffset(ctx, offset) + outputOffset, err := u.addWeightStringToOffset(ctx, offset) + if err != nil { + panic(err) + } + return outputOffset } func (u *Union) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gb bool, expr *sqlparser.AliasedExpr) int { @@ -193,7 +197,11 @@ func (u *Union) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gb bool, panic(vterrors.VT13001(fmt.Sprintf("could not find the argument to the weight_string function: %s", sqlparser.String(wsArg)))) } - return u.addWeightStringToOffset(ctx, argIdx) + offset, err := u.addWeightStringToOffset(ctx, argIdx) + if err != nil { + panic(err) + } + return offset case *sqlparser.Literal, *sqlparser.Argument: return u.addConstantToUnion(ctx, expr) default: @@ -278,7 +286,7 @@ func (u *Union) addConstantToUnion(ctx *plancontext.PlanningContext, aexpr *sqlp return } -func (u *Union) addWeightStringToOffset(ctx *plancontext.PlanningContext, argIdx int) (outputOffset int) { +func (u *Union) addWeightStringToOffset(ctx *plancontext.PlanningContext, argIdx int) (outputOffset int, err error) { for i, src := range u.Sources { thisOffset := src.AddWSColumn(ctx, argIdx, false) @@ -287,7 +295,7 @@ func (u *Union) addWeightStringToOffset(ctx *plancontext.PlanningContext, argIdx outputOffset = thisOffset } else { if thisOffset != outputOffset { - panic(vterrors.VT13001("weight_string offsets did not line up for UNION")) + return 0, vterrors.VT13001("weight_string offsets did not line up for UNION") } } } diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json index 7d02e4c401a..a92f476cc25 100644 --- a/go/vt/vtgate/planbuilder/testdata/cte_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/cte_cases.json @@ -1815,8 +1815,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, dt.c1 as foo, weight_string(dt.c0), weight_string(dt.c1) from (select id, foo from (select id, foo from `user` where 1 != 1) as x where 1 != 1 union select id, foo from (select id, foo from `user` where 1 != 1) as x where 1 != 1) as dt(c0, c1) where 1 != 1", - "Query": "select dt.c0 as id, dt.c1 as foo, weight_string(dt.c0), weight_string(dt.c1) from (select id, foo from (select id, foo from `user`) as x union select id, foo from (select id, foo from `user`) as x) as dt(c0, c1)" + "FieldQuery": "select id, foo, weight_string(id), weight_string(foo) from (select id, foo from `user` where 1 != 1) as x where 1 != 1 union select id, foo, weight_string(id), weight_string(foo) from (select id, foo from `user` where 1 != 1) as x where 1 != 1", + "Query": "select id, foo, weight_string(id), weight_string(foo) from (select id, foo from `user`) as x union select id, foo, weight_string(id), weight_string(foo) from (select id, foo from `user`) as x" } ] }, diff --git a/go/vt/vtgate/planbuilder/testdata/large_union_cases.json b/go/vt/vtgate/planbuilder/testdata/large_union_cases.json index 7dcfff4378e..1bf5a62908a 100644 --- a/go/vt/vtgate/planbuilder/testdata/large_union_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/large_union_cases.json @@ -24,8 +24,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as content, dt.c1 as user_id, weight_string(dt.c0), weight_string(dt.c1) from ((select content, user_id from music where 1 != 1) union (select content, user_id from music where 1 != 1)) as dt(c0, c1) where 1 != 1", - "Query": "select dt.c0 as content, dt.c1 as user_id, weight_string(dt.c0), weight_string(dt.c1) from ((select content, user_id from music where user_id = 1270698330 order by created_at asc, id asc limit 11) union (select content, user_id from music where user_id = 1270698330 order by created_at asc, id asc limit 11)) as dt(c0, c1)", + "FieldQuery": "(select content, user_id, weight_string(content), weight_string(user_id) from music where 1 != 1) union (select content, user_id, weight_string(content), weight_string(user_id) from music where 1 != 1)", + "Query": "(select content, user_id, weight_string(content), weight_string(user_id) from music where user_id = 1270698330 order by created_at asc, id asc limit 11) union (select content, user_id, weight_string(content), weight_string(user_id) from music where user_id = 1270698330 order by created_at asc, id asc limit 11)", "Values": [ "1270698330" ], @@ -38,8 +38,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as content, dt.c1 as user_id, weight_string(dt.c0), weight_string(dt.c1) from ((select content, user_id from music where 1 != 1) union (select content, user_id from music where 1 != 1)) as dt(c0, c1) where 1 != 1", - "Query": "select dt.c0 as content, dt.c1 as user_id, weight_string(dt.c0), weight_string(dt.c1) from ((select content, user_id from music where user_id = 1270699497 order by created_at asc, id asc limit 11) union (select content, user_id from music where user_id = 1270699497 order by created_at asc, id asc limit 11)) as dt(c0, c1)", + "FieldQuery": "(select content, user_id, weight_string(content), weight_string(user_id) from music where 1 != 1) union (select content, user_id, weight_string(content), weight_string(user_id) from music where 1 != 1)", + "Query": "(select content, user_id, weight_string(content), weight_string(user_id) from music where user_id = 1270699497 order by created_at asc, id asc limit 11) union (select content, user_id, weight_string(content), weight_string(user_id) from music where user_id = 1270699497 order by created_at asc, id asc limit 11)", "Values": [ "1270699497" ], diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json index 7fdad1b78d0..e8b19f93cad 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json @@ -1700,8 +1700,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from music where 1 != 1 union select id from music where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from music where music.user_id in ::__vals union select id from music where music.user_id in (1, 2, 3)) as dt(c0)", + "FieldQuery": "select id, weight_string(id) from music where 1 != 1 union select id, weight_string(id) from music where 1 != 1", + "Query": "select id, weight_string(id) from music where music.user_id in ::__vals union select id, weight_string(id) from music where music.user_id in (1, 2, 3)", "Values": [ "(1, 2, 3)" ], @@ -1735,8 +1735,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from music where 1 != 1 union select id from music where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from music where music.user_id in ::__vals union select id from music where music.user_id in ::vtg1) as dt(c0)", + "FieldQuery": "select id, weight_string(id) from music where 1 != 1 union select id, weight_string(id) from music where 1 != 1", + "Query": "select id, weight_string(id) from music where music.user_id in ::__vals union select id, weight_string(id) from music where music.user_id in ::vtg1", "Values": [ "::vtg1" ], diff --git a/go/vt/vtgate/planbuilder/testdata/union_cases.json b/go/vt/vtgate/planbuilder/testdata/union_cases.json index 3a74b182b06..8b47c46aef1 100644 --- a/go/vt/vtgate/planbuilder/testdata/union_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/union_cases.json @@ -43,8 +43,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` where 1 != 1 union select id from music where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` union select id from music) as dt(c0)" + "FieldQuery": "select id, weight_string(id) from `user` where 1 != 1 union select id, weight_string(id) from music where 1 != 1", + "Query": "select id, weight_string(id) from `user` union select id, weight_string(id) from music" } ] }, @@ -608,8 +608,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from ((select id from `user` where 1 != 1) union (select id from `user` where 1 != 1)) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from ((select id from `user` order by id desc) union (select id from `user` order by id asc)) as dt(c0)" + "FieldQuery": "(select id, weight_string(id) from `user` where 1 != 1) union (select id, weight_string(id) from `user` where 1 != 1)", + "Query": "(select id, weight_string(id) from `user` order by id desc) union (select id, weight_string(id) from `user` order by id asc)" } ] }, @@ -860,8 +860,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` where 1 != 1 union select id + 1 from `user` where 1 != 1 union select user_id from user_extra where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` union select id + 1 from `user` union select user_id from user_extra) as dt(c0)" + "FieldQuery": "select id, weight_string(id) from `user` where 1 != 1 union select id + 1, weight_string(id + 1) from `user` where 1 != 1 union select user_id, weight_string(user_id) from user_extra where 1 != 1", + "Query": "select id, weight_string(id) from `user` union select id + 1, weight_string(id + 1) from `user` union select user_id, weight_string(user_id) from user_extra" } ] }, @@ -895,8 +895,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` where 1 != 1 union select id from music where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` union select id from music) as dt(c0)" + "FieldQuery": "select id, weight_string(id) from `user` where 1 != 1 union select id, weight_string(id) from music where 1 != 1", + "Query": "select id, weight_string(id) from `user` union select id, weight_string(id) from music" }, { "OperatorType": "Route", @@ -1044,8 +1044,8 @@ "Name": "main", "Sharded": false }, - "FieldQuery": "select dt.c0 as col, weight_string(dt.c0) from (select col from unsharded where 1 != 1 union select col2 from unsharded where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as col, weight_string(dt.c0) from (select col from unsharded union select col2 from unsharded) as dt(c0)" + "FieldQuery": "select col, weight_string(col) from unsharded where 1 != 1 union select col2, weight_string(col2) from unsharded where 1 != 1", + "Query": "select col, weight_string(col) from unsharded union select col2, weight_string(col2) from unsharded" }, { "OperatorType": "Route", @@ -1054,8 +1054,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` where 1 != 1 union select col from user_extra where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` union select col from user_extra) as dt(c0)" + "FieldQuery": "select id, weight_string(id) from `user` where 1 != 1 union select col, weight_string(col) from user_extra where 1 != 1", + "Query": "select id, weight_string(id) from `user` union select col, weight_string(col) from user_extra" } ] } @@ -1190,8 +1190,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` where 1 != 1 union select 3 from dual where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as id, weight_string(dt.c0) from (select id from `user` union select 3 from dual) as dt(c0)" + "FieldQuery": "select id, weight_string(id) from `user` where 1 != 1 union select 3, weight_string(3) from dual where 1 != 1", + "Query": "select id, weight_string(id) from `user` union select 3, weight_string(3) from dual" } ] } @@ -1219,8 +1219,8 @@ { "OperatorType": "Distinct", "Collations": [ - "(0:2)", - "(1:3)", + "(0:3)", + "(1:4)", "2" ], "Inputs": [ @@ -1231,8 +1231,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as id, dt.c1 as id, weight_string(dt.c0), weight_string(dt.c1) from (select id, id from `user` where 1 != 1 union select id, bar from user_extra where 1 != 1) as dt(c0, c1) where 1 != 1", - "Query": "select dt.c0 as id, dt.c1 as id, weight_string(dt.c0), weight_string(dt.c1) from (select id, id from `user` union select id, bar from user_extra) as dt(c0, c1)" + "FieldQuery": "select id, id, weight_string(id), weight_string(id), weight_string(id) from `user` where 1 != 1 union select id, bar, weight_string(id), weight_string(id), weight_string(bar) from user_extra where 1 != 1", + "Query": "select id, id, weight_string(id), weight_string(id), weight_string(id) from `user` union select id, bar, weight_string(id), weight_string(id), weight_string(bar) from user_extra" } ] } @@ -1598,8 +1598,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as bar, dt.c1 as baz, dt.c2 as toto, weight_string(dt.c0), weight_string(dt.c1), weight_string(dt.c2) from (select bar, baz, toto from music where 1 != 1 union select foo, foo, foo from `user` where 1 != 1) as dt(c0, c1, c2) where 1 != 1", - "Query": "select dt.c0 as bar, dt.c1 as baz, dt.c2 as toto, weight_string(dt.c0), weight_string(dt.c1), weight_string(dt.c2) from (select bar, baz, toto from music union select foo, foo, foo from `user`) as dt(c0, c1, c2)" + "FieldQuery": "select bar, baz, toto, weight_string(bar), weight_string(baz), weight_string(toto) from music where 1 != 1 union select foo, foo, foo, weight_string(foo), weight_string(foo), weight_string(foo) from `user` where 1 != 1", + "Query": "select bar, baz, toto, weight_string(bar), weight_string(baz), weight_string(toto) from music union select foo, foo, foo, weight_string(foo), weight_string(foo), weight_string(foo) from `user`" } ] }, @@ -1632,8 +1632,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as foo, dt.c1 as foo, dt.c2 as foo, weight_string(dt.c0), weight_string(dt.c1), weight_string(dt.c2) from (select foo, foo, foo from `user` where 1 != 1 union select bar, baz, toto from music where 1 != 1) as dt(c0, c1, c2) where 1 != 1", - "Query": "select dt.c0 as foo, dt.c1 as foo, dt.c2 as foo, weight_string(dt.c0), weight_string(dt.c1), weight_string(dt.c2) from (select foo, foo, foo from `user` union select bar, baz, toto from music) as dt(c0, c1, c2)" + "FieldQuery": "select foo, foo, foo, weight_string(foo), weight_string(foo), weight_string(foo) from `user` where 1 != 1 union select bar, baz, toto, weight_string(bar), weight_string(baz), weight_string(toto) from music where 1 != 1", + "Query": "select foo, foo, foo, weight_string(foo), weight_string(foo), weight_string(foo) from `user` union select bar, baz, toto, weight_string(bar), weight_string(baz), weight_string(toto) from music" } ] }, @@ -1671,8 +1671,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as foo, weight_string(dt.c0) from (select foo from `user` where 1 != 1 union select foo from `user` where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as foo, weight_string(dt.c0) from (select foo from `user` where bar = 12 union select foo from `user` where bar = 134) as dt(c0)" + "FieldQuery": "select foo, weight_string(foo) from `user` where 1 != 1 union select foo, weight_string(foo) from `user` where 1 != 1", + "Query": "select foo, weight_string(foo) from `user` where bar = 12 union select foo, weight_string(foo) from `user` where bar = 134" } ] }, @@ -1689,8 +1689,8 @@ "Name": "user", "Sharded": true }, - "FieldQuery": "select dt.c0 as bar, weight_string(dt.c0) from (select bar from music where 1 != 1 union select bar from music where 1 != 1) as dt(c0) where 1 != 1", - "Query": "select dt.c0 as bar, weight_string(dt.c0) from (select bar from music where foo = 12 and bar = :t1_foo union select bar from music where foo = 1234 and bar = :t1_foo) as dt(c0)" + "FieldQuery": "select bar, weight_string(bar) from music where 1 != 1 union select bar, weight_string(bar) from music where 1 != 1", + "Query": "select bar, weight_string(bar) from music where foo = 12 and bar = :t1_foo union select bar, weight_string(bar) from music where foo = 1234 and bar = :t1_foo" } ] }