Skip to content

Commit

Permalink
Fix CTE query by fixing bindvar pushing in unions (#15838)
Browse files Browse the repository at this point in the history 
Signed-off-by: Manan Gupta <[email protected]>
  • Loading branch information
@GuptaManan100
GuptaManan100 authored May 6, 2024
1 parent f27d287 commit 94259cd
Show file tree
Hide file tree
Showing 4 changed files with 177 additions and 0 deletions.
24 changes: 24 additions & 0 deletions go/test/endtoend/vtgate/queries/derived/cte_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,27 @@ func TestCTEColumns(t *testing.T) {
mcmp.AssertMatches(`with t(id) as (SELECT id FROM user) SELECT t.id FROM t ORDER BY t.id DESC`,
`[[INT64(5)] [INT64(4)] [INT64(3)] [INT64(2)] [INT64(1)]]`)
}

func TestCTEAggregationsInUnion(t *testing.T) {
utils.SkipIfBinaryIsBelowVersion(t, 20, "vtgate")
mcmp, closer := start(t)
defer closer()

mcmp.AssertMatches(`WITH toto AS (SELECT COUNT(*) as num
FROM (SELECT user.id
FROM user
WHERE user.name = 'toto'
LIMIT 1000) t LIMIT 1 ),
tata AS (SELECT COUNT(*) as num
FROM (SELECT user.id
FROM user
WHERE user.name = 'tata'
LIMIT 1000) t LIMIT 1),
total AS (SELECT LEAST(1000, SUM(num)) AS num
FROM (SELECT num
FROM toto
UNION ALL SELECT num
FROM tata) t LIMIT 1)
SELECT 'total' AS tab, num
FROM total`, `[[VARCHAR("total") DECIMAL(2)]]`)
}
3 changes: 3 additions & 0 deletions go/vt/vtgate/planbuilder/operators/horizon_expanding.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,9 @@ outer:
func createProjectionForComplexAggregation(a *Aggregator, qp *QueryProjection) Operator {
p := newAliasedProjection(a)
p.DT = a.DT
// We don't want to keep the derived table in both Aggregator and Projection.
// If we do, then we end up re-writing the same column twice.
a.DT = nil
for _, expr := range qp.SelectExprs {
ae, err := expr.GetAliasedExpr()
if err != nil {
Expand Down
17 changes: 17 additions & 0 deletions go/vt/vtgate/planbuilder/operators/union.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,11 +193,28 @@ func (u *Union) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gb bool,
}

return u.addWeightStringToOffset(ctx, argIdx)
case *sqlparser.Literal, *sqlparser.Argument:
return u.addConstantToUnion(ctx, expr)
default:
panic(vterrors.VT13001(fmt.Sprintf("only weight_string function is expected - got %s", sqlparser.String(expr))))
}
}

func (u *Union) addConstantToUnion(ctx *plancontext.PlanningContext, aexpr *sqlparser.AliasedExpr) (outputOffset int) {
for i, src := range u.Sources {
thisOffset := src.AddColumn(ctx, true, false, aexpr)
// all offsets for the newly added ws need to line up
if i == 0 {
outputOffset = thisOffset
} else {
if thisOffset != outputOffset {
panic(vterrors.VT12001("argument offsets did not line up for UNION"))
}
}
}
return
}

func (u *Union) addWeightStringToOffset(ctx *plancontext.PlanningContext, argIdx int) (outputOffset int) {
for i, src := range u.Sources {
thisOffset := src.AddWSColumn(ctx, argIdx, false)
Expand Down
133 changes: 133 additions & 0 deletions go/vt/vtgate/planbuilder/testdata/cte_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1957,5 +1957,138 @@
"user.user_metadata"
]
}
},
{
"comment": "CTE expression using unions and complex aggregation with literal",
"query": "WITH `open` AS (SELECT COUNT(*) as `num` FROM (SELECT `user`.`id` FROM `user` WHERE `user`.`textcol1` = 'open' AND `user`.`intcol` = 1 LIMIT 1000) `t` LIMIT 1 ), `closed` AS (SELECT COUNT(*) as `num` FROM ( SELECT `user`.`id` FROM `user` WHERE `user`.`textcol1` = 'closed' AND `user`.`intcol` = 1 LIMIT 1000) `t` LIMIT 1 ), `all` AS (SELECT LEAST(1000, SUM(`num`)) AS `num` FROM ( SELECT `num` FROM `open` UNION ALL SELECT `num` FROM `closed` ) `t` LIMIT 1 )SELECT 'all' AS `tab`, `num`FROM `all`",
"plan": {
"QueryType": "SELECT",
"Original": "WITH `open` AS (SELECT COUNT(*) as `num` FROM (SELECT `user`.`id` FROM `user` WHERE `user`.`textcol1` = 'open' AND `user`.`intcol` = 1 LIMIT 1000) `t` LIMIT 1 ), `closed` AS (SELECT COUNT(*) as `num` FROM ( SELECT `user`.`id` FROM `user` WHERE `user`.`textcol1` = 'closed' AND `user`.`intcol` = 1 LIMIT 1000) `t` LIMIT 1 ), `all` AS (SELECT LEAST(1000, SUM(`num`)) AS `num` FROM ( SELECT `num` FROM `open` UNION ALL SELECT `num` FROM `closed` ) `t` LIMIT 1 )SELECT 'all' AS `tab`, `num`FROM `all`",
"Instructions": {
"OperatorType": "Limit",
"Count": "1",
"Inputs": [
{
"OperatorType": "Projection",
"Expressions": [
"'all' as tab",
":0 as num"
],
"Inputs": [
{
"OperatorType": "Projection",
"Expressions": [
"least(1000, sum(num)) as num"
],
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Scalar",
"Aggregates": "any_value(0), sum(1) AS sum(num)",
"Inputs": [
{
"OperatorType": "SimpleProjection",
"Columns": [
1,
0
],
"Inputs": [
{
"OperatorType": "Concatenate",
"Inputs": [
{
"OperatorType": "Limit",
"Count": "1",
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Scalar",
"Aggregates": "count_star(0) AS num, any_value(1)",
"Inputs": [
{
"OperatorType": "SimpleProjection",
"Columns": [
1,
2
],
"Inputs": [
{
"OperatorType": "Limit",
"Count": "1000",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select t.id, 1, 1000 from (select `user`.id from `user` where 1 != 1) as t where 1 != 1",
"Query": "select t.id, 1, 1000 from (select `user`.id from `user` where `user`.textcol1 = 'open' and `user`.intcol = 1) as t limit :__upper_limit",
"Table": "`user`"
}
]
}
]
}
]
}
]
},
{
"OperatorType": "Limit",
"Count": "1",
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Scalar",
"Aggregates": "count_star(0) AS num, any_value(1)",
"Inputs": [
{
"OperatorType": "SimpleProjection",
"Columns": [
1,
2
],
"Inputs": [
{
"OperatorType": "Limit",
"Count": "1000",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select t.id, 1, 1000 from (select `user`.id from `user` where 1 != 1) as t where 1 != 1",
"Query": "select t.id, 1, 1000 from (select `user`.id from `user` where `user`.textcol1 = 'closed' and `user`.intcol = 1) as t limit :__upper_limit",
"Table": "`user`"
}
]
}
]
}
]
}
]
}
]
}
]
}
]
}
]
}
]
}
]
},
"TablesUsed": [
"user.user"
]
}
}
]

0 comments on commit 94259cd

Please sign in to comment.