Skip to content

Commit

Permalink
Add support for more queries (#14369)
Browse files Browse the repository at this point in the history
Co-authored-by: Manan Gupta <[email protected]>
  • Loading branch information
systay and GuptaManan100 authored Nov 1, 2023
1 parent 3924566 commit b26a2d1
Show file tree
Hide file tree
Showing 5 changed files with 192 additions and 121 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -206,10 +206,6 @@ func TestMultipleSchemaPredicates(t *testing.T) {
}

func TestInfrSchemaAndUnionAll(t *testing.T) {
clusterInstance.VtGateExtraArgs = append(clusterInstance.VtGateExtraArgs, "--planner-version=gen4")
require.NoError(t,
clusterInstance.RestartVtgate())

vtConnParams := clusterInstance.GetVTParams(keyspaceName)
vtConnParams.DbName = keyspaceName
conn, err := mysql.Connect(context.Background(), &vtConnParams)
Expand All @@ -225,3 +221,39 @@ func TestInfrSchemaAndUnionAll(t *testing.T) {
})
}
}

func TestTypeORMQuery(t *testing.T) {
// This test checks that we can run queries similar to the ones that the TypeORM framework uses

require.NoError(t,
utils.WaitForAuthoritative(t, "ks", "t1", clusterInstance.VtgateProcess.ReadVSchema))

mcmp, closer := start(t)
defer closer()

query := `SELECT kcu.TABLE_NAME, kcu.COLUMN_NAME, cols.DATA_TYPE
FROM (SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME
FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu
WHERE kcu.TABLE_SCHEMA = 'ks'
AND kcu.TABLE_NAME = 't1'
UNION
SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME
FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu
WHERE kcu.TABLE_SCHEMA = 'ks'
AND kcu.TABLE_NAME = 't7_xxhash') kcu
INNER JOIN (SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, DATA_TYPE
FROM INFORMATION_SCHEMA.COLUMNS cols
WHERE cols.TABLE_SCHEMA = 'ks'
AND cols.TABLE_NAME = 't1'
UNION
SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, DATA_TYPE
FROM INFORMATION_SCHEMA.COLUMNS cols
WHERE cols.TABLE_SCHEMA = 'ks'
AND cols.TABLE_NAME = 't7_xxhash') cols
ON kcu.TABLE_SCHEMA = cols.TABLE_SCHEMA AND kcu.TABLE_NAME = cols.TABLE_NAME AND
kcu.COLUMN_NAME = cols.COLUMN_NAME`
utils.AssertMatchesAny(t, mcmp.VtConn, query,
`[[VARBINARY("t1") VARCHAR("id1") BLOB("bigint")] [VARBINARY("t7_xxhash") VARCHAR("uid") BLOB("varchar")]]`,
`[[VARCHAR("t1") VARCHAR("id1") BLOB("bigint")] [VARCHAR("t7_xxhash") VARCHAR("uid") BLOB("varchar")]]`,
)
}
4 changes: 0 additions & 4 deletions go/vt/vtgate/planbuilder/operators/projection.go
Original file line number Diff line number Diff line change
Expand Up @@ -455,10 +455,6 @@ func (p *Projection) ShortDescription() string {
}

func (p *Projection) Compact(ctx *plancontext.PlanningContext) (ops.Operator, *rewrite.ApplyResult, error) {
if p.isDerived() {
return p, rewrite.SameTree, nil
}

ap, err := p.GetAliasedProjections()
if err != nil {
return p, rewrite.SameTree, nil
Expand Down
2 changes: 1 addition & 1 deletion go/vt/vtgate/planbuilder/operators/route_planning.go
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,7 @@ func requiresSwitchingSides(ctx *plancontext.PlanningContext, op ops.Operator) b
_ = rewrite.Visit(op, func(current ops.Operator) error {
horizon, isHorizon := current.(*Horizon)

if isHorizon && horizon.IsDerived() && !horizon.IsMergeable(ctx) {
if isHorizon && !horizon.IsMergeable(ctx) {
required = true
return io.EOF
}
Expand Down
49 changes: 26 additions & 23 deletions go/vt/vtgate/planbuilder/operators/subquery_planning.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ limitations under the License.
package operators

import (
"fmt"
"io"

"golang.org/x/exp/slices"
Expand All @@ -41,34 +42,36 @@ func isMergeable(ctx *plancontext.PlanningContext, query sqlparser.SelectStateme
return false
}

sel, ok := query.(*sqlparser.Select)
if !ok {
return false
}

if len(sel.GroupBy) > 0 {
// iff we are grouping, we need to check that we can perform the grouping inside a single shard, and we check that
// by checking that one of the grouping expressions used is a unique single column vindex.
// TODO: we could also support the case where all the columns of a multi-column vindex are used in the grouping
for _, gb := range sel.GroupBy {
if validVindex(gb) {
return true
switch node := query.(type) {
case *sqlparser.Select:
if len(node.GroupBy) > 0 {
// iff we are grouping, we need to check that we can perform the grouping inside a single shard, and we check that
// by checking that one of the grouping expressions used is a unique single column vindex.
// TODO: we could also support the case where all the columns of a multi-column vindex are used in the grouping
for _, gb := range node.GroupBy {
if validVindex(gb) {
return true
}
}
return false
}
return false
}

// if we have grouping, we have already checked that it's safe, and don't need to check for aggregations
// but if we don't have groupings, we need to check if there are aggregations that will mess with us
if sqlparser.ContainsAggregation(sel.SelectExprs) {
return false
}
// if we have grouping, we have already checked that it's safe, and don't need to check for aggregations
// but if we don't have groupings, we need to check if there are aggregations that will mess with us
if sqlparser.ContainsAggregation(node.SelectExprs) {
return false
}

if sqlparser.ContainsAggregation(sel.Having) {
return false
}
if sqlparser.ContainsAggregation(node.Having) {
return false
}

return true
return true
case *sqlparser.Union:
return isMergeable(ctx, node.Left, op) && isMergeable(ctx, node.Right, op)
default:
panic(vterrors.VT13001(fmt.Sprintf("Unknown SelectStatement type - %T", node)))
}
}

func settleSubqueries(ctx *plancontext.PlanningContext, op ops.Operator) ops.Operator {
Expand Down
218 changes: 129 additions & 89 deletions go/vt/vtgate/planbuilder/testdata/union_cases.json
Original file line number Diff line number Diff line change
Expand Up @@ -749,29 +749,22 @@
"QueryType": "SELECT",
"Original": "select * from ((select id from user union select id+1 from user) union select user_id from user_extra) as t",
"Instructions": {
"OperatorType": "SimpleProjection",
"Columns": [
0
"OperatorType": "Distinct",
"Collations": [
"(0:1)"
],
"ResultColumns": 1,
"Inputs": [
{
"OperatorType": "Distinct",
"Collations": [
"(0:1)"
],
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, weight_string(id) 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 where 1 != 1",
"Query": "select id, weight_string(id) from (select id from `user` union select id + 1 from `user` union select user_id from user_extra) as dt",
"Table": "`user`, user_extra"
}
]
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, weight_string(id) 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 where 1 != 1",
"Query": "select id, weight_string(id) from (select id from `user` union select id + 1 from `user` union select user_id from user_extra) as dt",
"Table": "`user`, user_extra"
}
]
},
Expand Down Expand Up @@ -955,49 +948,41 @@
"TableName": "`user`_`user`",
"Inputs": [
{
"OperatorType": "SimpleProjection",
"Columns": [
0
],
"OperatorType": "Concatenate",
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "Limit",
"Count": "5",
"Inputs": [
{
"OperatorType": "Limit",
"Count": "5",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, weight_string(id) from `user` where 1 != 1",
"OrderBy": "(0|1) ASC",
"Query": "select id, weight_string(id) from `user` order by id asc limit :__upper_limit",
"Table": "`user`"
}
]
},
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, weight_string(id) from `user` where 1 != 1",
"OrderBy": "(0|1) ASC",
"Query": "select id, weight_string(id) from `user` order by id asc limit :__upper_limit",
"Table": "`user`"
}
]
},
{
"OperatorType": "Limit",
"Count": "5",
"Inputs": [
{
"OperatorType": "Limit",
"Count": "5",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, weight_string(id) from `user` where 1 != 1",
"OrderBy": "(0|1) DESC",
"Query": "select id, weight_string(id) from `user` order by id desc limit :__upper_limit",
"Table": "`user`"
}
]
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, weight_string(id) from `user` where 1 != 1",
"OrderBy": "(0|1) DESC",
"Query": "select id, weight_string(id) from `user` order by id desc limit :__upper_limit",
"Table": "`user`"
}
]
}
Expand Down Expand Up @@ -1098,41 +1083,35 @@
],
"Inputs": [
{
"OperatorType": "Projection",
"Expressions": null,
"OperatorType": "Distinct",
"Collations": [
"(0:1)"
],
"Inputs": [
{
"OperatorType": "Distinct",
"Collations": [
"(0:1)"
],
"OperatorType": "Concatenate",
"Inputs": [
{
"OperatorType": "Concatenate",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id + 42 as foo, weight_string(id + 42) from `user` where 1 != 1",
"Query": "select distinct id + 42 as foo, weight_string(id + 42) from `user`",
"Table": "`user`"
},
{
"OperatorType": "Route",
"Variant": "Unsharded",
"Keyspace": {
"Name": "main",
"Sharded": false
},
"FieldQuery": "select 1 + id as foo, weight_string(1 + id) from unsharded where 1 != 1",
"Query": "select distinct 1 + id as foo, weight_string(1 + id) from unsharded",
"Table": "unsharded"
}
]
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id + 42 as foo, weight_string(id + 42) from `user` where 1 != 1",
"Query": "select distinct id + 42 as foo, weight_string(id + 42) from `user`",
"Table": "`user`"
},
{
"OperatorType": "Route",
"Variant": "Unsharded",
"Keyspace": {
"Name": "main",
"Sharded": false
},
"FieldQuery": "select 1 + id as foo, weight_string(1 + id) from unsharded where 1 != 1",
"Query": "select distinct 1 + id as foo, weight_string(1 + id) from unsharded",
"Table": "unsharded"
}
]
}
Expand Down Expand Up @@ -1495,5 +1474,66 @@
"user.user"
]
}
},
{
"comment": "join between two derived tables containing UNION",
"query": "select * from (select foo from user where bar = 12 union select foo from user where bar = 134) as t1 join (select bar from music where foo = 12 union select bar from music where foo = 1234) as t2 on t1.foo = t2.bar",
"plan": {
"QueryType": "SELECT",
"Original": "select * from (select foo from user where bar = 12 union select foo from user where bar = 134) as t1 join (select bar from music where foo = 12 union select bar from music where foo = 1234) as t2 on t1.foo = t2.bar",
"Instructions": {
"OperatorType": "Join",
"Variant": "Join",
"JoinColumnIndexes": "L:0,R:0",
"JoinVars": {
"t1_foo": 0
},
"TableName": "`user`_music",
"Inputs": [
{
"OperatorType": "Distinct",
"Collations": [
"(0:1)"
],
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select foo, weight_string(foo) from (select foo from `user` where 1 != 1 union select foo from `user` where 1 != 1) as dt where 1 != 1",
"Query": "select foo, weight_string(foo) from (select foo from `user` where bar = 12 union select foo from `user` where bar = 134) as dt",
"Table": "`user`"
}
]
},
{
"OperatorType": "Distinct",
"Collations": [
"(0:1)"
],
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select bar, weight_string(bar) from (select bar from music where 1 != 1 union select bar from music where 1 != 1) as dt where 1 != 1",
"Query": "select bar, weight_string(bar) 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",
"Table": "music"
}
]
}
]
},
"TablesUsed": [
"user.music",
"user.user"
]
}
}
]

0 comments on commit b26a2d1

Please sign in to comment.