Skip to content

Commit 4d05a25

Browse files
arthurschreiberarthurschreiber
committed
Better support for merging UNION routes that use IN conditions.
Signed-off-by: Arthur Schreiber <[email protected]>
1 parent 1140923 commit 4d05a25

File tree

3 files changed

+116
-6
lines changed

3 files changed

+116
-6
lines changed

‎go/vt/vtgate/planbuilder/operators/route_planning.go

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -478,14 +478,26 @@ func gen4ValuesEqual(ctx *plancontext.PlanningContext, a, b []sqlparser.Expr) (b
478478
return false, nil
479479
}
480480
if c != nil {
481-
conditions = append(conditions, *c)
481+
conditions = append(conditions, c...)
482482
}
483483
}
484484
return true, conditions
485485
}
486486

487-
func gen4ValEqual(ctx *plancontext.PlanningContext, a, b sqlparser.Expr) (bool, *engine.Condition) {
487+
func gen4ValEqual(ctx *plancontext.PlanningContext, a, b sqlparser.Expr) (bool, []engine.Condition) {
488488
switch a := a.(type) {
489+
case sqlparser.ValTuple:
490+
if b, ok := b.(sqlparser.ValTuple); ok {
491+
return gen4ValuesEqual(ctx, a, b)
492+
}
493+
494+
return false, nil
495+
496+
case sqlparser.ListArg:
497+
if b, ok := b.(sqlparser.ListArg); ok {
498+
return a == b, nil
499+
}
500+
489501
case *sqlparser.ColName:
490502
if b, ok := b.(*sqlparser.ColName); ok {
491503
if !a.Name.Equal(b.Name) {
@@ -518,7 +530,7 @@ func gen4ValEqual(ctx *plancontext.PlanningContext, a, b sqlparser.Expr) (bool,
518530
}
519531

520532
return aVal.Type == bVal.Type && bytes.Equal(aVal.Value, bVal.Value),
521-
&engine.Condition{A: a.Name, B: b.Name}
533+
[]engine.Condition{{A: a.Name, B: b.Name}}
522534

523535
case *sqlparser.Literal:
524536
b, ok := b.(*sqlparser.Literal)

‎go/vt/vtgate/planbuilder/operators/union_merging.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,8 +146,6 @@ func tryMergeUnionShardedRouting(
146146

147147
scatterA := tblA.RouteOpCode == engine.Scatter
148148
scatterB := tblB.RouteOpCode == engine.Scatter
149-
uniqueA := tblA.RouteOpCode == engine.EqualUnique
150-
uniqueB := tblB.RouteOpCode == engine.EqualUnique
151149

152150
switch {
153151
case scatterA:
@@ -156,7 +154,11 @@ func tryMergeUnionShardedRouting(
156154
case scatterB:
157155
return createMergedUnion(ctx, routeA, routeB, exprsA, exprsB, distinct, tblB, nil)
158156

159-
case uniqueA && uniqueB:
157+
case tblA.RouteOpCode == engine.EqualUnique && tblB.RouteOpCode == engine.EqualUnique:
158+
fallthrough
159+
case tblA.RouteOpCode == engine.Equal && tblB.RouteOpCode == engine.Equal:
160+
fallthrough
161+
case tblA.RouteOpCode == engine.IN && tblB.RouteOpCode == engine.IN:
160162
aVdx := tblA.SelectedVindex()
161163
bVdx := tblB.SelectedVindex()
162164
aExpr := tblA.VindexExpressions()

‎go/vt/vtgate/planbuilder/testdata/select_cases.json

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1679,6 +1679,102 @@
16791679
},
16801680
"skip_e2e": true
16811681
},
1682+
{
1683+
"comment": "union with the same target shard (using `IN` with multiple matching values)",
1684+
"query": "SELECT id FROM music WHERE music.user_id IN (1, 2, 3) UNION SELECT id FROM music WHERE music.user_id IN (1, 2, 3)",
1685+
"plan": {
1686+
"Type": "Complex",
1687+
"QueryType": "SELECT",
1688+
"Original": "SELECT id FROM music WHERE music.user_id IN (1, 2, 3) UNION SELECT id FROM music WHERE music.user_id IN (1, 2, 3)",
1689+
"Instructions": {
1690+
"OperatorType": "Distinct",
1691+
"Collations": [
1692+
"(0:1)"
1693+
],
1694+
"ResultColumns": 1,
1695+
"Inputs": [
1696+
{
1697+
"OperatorType": "Route",
1698+
"Variant": "IN",
1699+
"Keyspace": {
1700+
"Name": "user",
1701+
"Sharded": true
1702+
},
1703+
"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",
1704+
"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)",
1705+
"Values": [
1706+
"(1, 2, 3)"
1707+
],
1708+
"Vindex": "user_index"
1709+
}
1710+
]
1711+
},
1712+
"TablesUsed": [
1713+
"user.music"
1714+
]
1715+
}
1716+
},
1717+
{
1718+
"comment": "union with the same target shard (using `IN` with the same normalized value)",
1719+
"query": "SELECT id FROM music WHERE music.user_id IN ::vtg1 UNION SELECT id FROM music WHERE music.user_id IN ::vtg1",
1720+
"plan": {
1721+
"Type": "Complex",
1722+
"QueryType": "SELECT",
1723+
"Original": "SELECT id FROM music WHERE music.user_id IN ::vtg1 UNION SELECT id FROM music WHERE music.user_id IN ::vtg1",
1724+
"Instructions": {
1725+
"OperatorType": "Distinct",
1726+
"Collations": [
1727+
"(0:1)"
1728+
],
1729+
"ResultColumns": 1,
1730+
"Inputs": [
1731+
{
1732+
"OperatorType": "Route",
1733+
"Variant": "IN",
1734+
"Keyspace": {
1735+
"Name": "user",
1736+
"Sharded": true
1737+
},
1738+
"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",
1739+
"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)",
1740+
"Values": [
1741+
"::vtg1"
1742+
],
1743+
"Vindex": "user_index"
1744+
}
1745+
]
1746+
},
1747+
"TablesUsed": [
1748+
"user.music"
1749+
]
1750+
}
1751+
},
1752+
{
1753+
"comment": "union with the same target shard (using a mix of `IN` and `=` conditions)",
1754+
"query": "SELECT id FROM music WHERE music.user_id IN (1) UNION SELECT id FROM music WHERE music.user_id = 1",
1755+
"plan": {
1756+
"Type": "Passthrough",
1757+
"QueryType": "SELECT",
1758+
"Original": "SELECT id FROM music WHERE music.user_id IN (1) UNION SELECT id FROM music WHERE music.user_id = 1",
1759+
"Instructions": {
1760+
"OperatorType": "Route",
1761+
"Variant": "EqualUnique",
1762+
"Keyspace": {
1763+
"Name": "user",
1764+
"Sharded": true
1765+
},
1766+
"FieldQuery": "select id from music where 1 != 1 union select id from music where 1 != 1",
1767+
"Query": "select id from music where music.user_id in (1) union select id from music where music.user_id = 1",
1768+
"Values": [
1769+
"1"
1770+
],
1771+
"Vindex": "user_index"
1772+
},
1773+
"TablesUsed": [
1774+
"user.music"
1775+
]
1776+
}
1777+
},
16821778
{
16831779
"comment": "union with the same target shard last_insert_id",
16841780
"query": "select *, last_insert_id() from music where user_id = 1 union select * from user where id = 1",

0 commit comments

Comments
 (0)