Skip to content

Commit

Permalink
planner: support union statements with ctes (#15312)
Browse files Browse the repository at this point in the history 
Signed-off-by: Max Englander <[email protected]>
  • Loading branch information
@maxenglander
maxenglander authored Feb 21, 2024
1 parent c982a04 commit a9b2c18
Show file tree
Hide file tree
Showing 4 changed files with 114 additions and 0 deletions.
4 changes: 4 additions & 0 deletions go/vt/sqlparser/ast_format.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,10 @@ func (node *CommentOnly) Format(buf *TrackedBuffer) {

// Format formats the node.
func (node *Union) Format(buf *TrackedBuffer) {
if node.With != nil {
buf.astPrintf(node, "%v", node.With)
}

if requiresParen(node.Left) {
buf.astPrintf(node, "(%v)", node.Left)
} else {
Expand Down
4 changes: 4 additions & 0 deletions go/vt/sqlparser/ast_format_fast.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions go/vt/sqlparser/parse_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -505,6 +505,9 @@ var (
}, {
input: "WITH topsales2003 AS (SELECT salesRepEmployeeNumber employeeNumber, SUM(quantityOrdered * priceEach) sales FROM orders INNER JOIN orderdetails USING (orderNumber) INNER JOIN customers USING (customerNumber) WHERE YEAR(shippedDate) = 2003 AND status = 'Shipped' GROUP BY salesRepEmployeeNumber ORDER BY sales DESC LIMIT 5)SELECT employeeNumber, firstName, lastName, sales FROM employees JOIN topsales2003 USING (employeeNumber)",
output: "with topsales2003 as (select salesRepEmployeeNumber as employeeNumber, sum(quantityOrdered * priceEach) as sales from orders join orderdetails using (orderNumber) join customers using (customerNumber) where YEAR(shippedDate) = 2003 and `status` = 'Shipped' group by salesRepEmployeeNumber order by sales desc limit 5) select employeeNumber, firstName, lastName, sales from employees join topsales2003 using (employeeNumber)",
}, {
input: "WITH count_a AS (SELECT COUNT(`id`) AS `num` FROM `tbl_a`), count_b AS (SELECT COUNT(`id`) AS `num` FROM tbl_b) SELECT 'a', `num` FROM `count_a` UNION SELECT 'b', `num` FROM `count_b`",
output: "with count_a as (select count(id) as num from tbl_a) , count_b as (select count(id) as num from tbl_b) select 'a', num from count_a union select 'b', num from count_b",
}, {
input: "select 1 from t",
}, {
Expand Down
103 changes: 103 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 @@ -1863,5 +1863,108 @@
"main.unsharded"
]
}
},
{
"comment": "WITH two common expressions against an unsharded datbase and a SELECT UNION against those expressions",
"query": "WITH `count_a` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_a`), `count_b` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_b`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`",
"plan": {
"QueryType": "SELECT",
"Original": "WITH `count_a` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_a`), `count_b` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_b`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`",
"Instructions": {
"OperatorType": "Route",
"Variant": "Unsharded",
"Keyspace": {
"Name": "main",
"Sharded": false
},
"FieldQuery": "select 'count_a' as tab, num from count_a where 1 != 1 union select 'count_b' as tab, num from count_b where 1 != 1",
"Query": "with count_a as (select count(id) as num from unsharded_a) , count_b as (select count(id) as num from unsharded_b) select 'count_a' as tab, num from count_a union select 'count_b' as tab, num from count_b",
"Table": "unsharded_a, unsharded_b"
},
"TablesUsed": [
"main.unsharded_a",
"main.unsharded_b"
]
}
},
{
"comment": "WITH two common expressions against a sharded datbase and a SELECT UNION against those expressions",
"query": "WITH `count_a` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_metadata`), `count_b` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_extra`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`",
"plan": {
"QueryType": "SELECT",
"Original": "WITH `count_a` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_metadata`), `count_b` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_extra`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`",
"Instructions": {
"OperatorType": "Distinct",
"Collations": [
"0: utf8mb4_0900_ai_ci",
"1"
],
"Inputs": [
{
"OperatorType": "Concatenate",
"Inputs": [
{
"OperatorType": "Projection",
"Expressions": [
"'count_a' as tab",
":0 as num"
],
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Scalar",
"Aggregates": "sum_count(0) AS num",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select count(user_id) as num from user_metadata where 1 != 1",
"Query": "select count(user_id) as num from user_metadata",
"Table": "user_metadata"
}
]
}
]
},
{
"OperatorType": "Projection",
"Expressions": [
"'count_b' as tab",
":0 as num"
],
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Scalar",
"Aggregates": "sum_count(0) AS num",
"Inputs": [
{
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select count(user_id) as num from user_extra where 1 != 1",
"Query": "select count(user_id) as num from user_extra",
"Table": "user_extra"
}
]
}
]
}
]
}
]
},
"TablesUsed": [
"user.user_extra",
"user.user_metadata"
]
}
}
]

0 comments on commit a9b2c18

Please sign in to comment.