Skip to content

Commit

Permalink
Stop truncating needed offsets
Browse files Browse the repository at this point in the history 
Once an offset has been shared with the outside,
we need to not truncate away that column

Signed-off-by: Florent Poinsard <[email protected]>
Signed-off-by: Andres Taylor <[email protected]>
  • Loading branch information
@frouioui @systay
frouioui authored and systay committed Jun 19, 2024
1 parent 6441298 commit 3039c4e
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 10 deletions.
10 changes: 10 additions & 0 deletions go/vt/vtgate/planbuilder/operators/aggregator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,11 +144,18 @@ func (a *Aggregator) FindCol(ctx *plancontext.PlanningContext, in sqlparser.Expr

expr := a.DT.RewriteExpression(ctx, in)
if offset, found := canReuseColumn(ctx, a.Columns, expr, extractExpr); found {
a.checkOffset(offset)
return offset
}
return -1
}

func (a *Aggregator) checkOffset(offset int) {
if a.ResultColumns > 0 && a.ResultColumns <= offset {
a.ResultColumns = offset + 1
}
}

func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, groupBy bool, ae *sqlparser.AliasedExpr) int {
rewritten := a.DT.RewriteExpression(ctx, ae.Expr)

Expand All @@ -160,6 +167,7 @@ func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gro
if reuse {
offset := a.findColInternal(ctx, ae, groupBy)
if offset >= 0 {
a.checkOffset(offset)
return offset
}
}
Expand Down Expand Up @@ -191,6 +199,7 @@ func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gro
panic(errFailedToPlan(ae))
}

a.checkOffset(offset)
return offset
}

Expand Down Expand Up @@ -249,6 +258,7 @@ func (a *Aggregator) AddWSColumn(ctx *plancontext.PlanningContext, offset int, u
// TODO: we could handle this case by adding a projection on under the aggregator to make the columns line up
panic(errFailedToPlan(wsAe))
}
a.checkOffset(wsOffset)
return wsOffset
}

Expand Down
4 changes: 2 additions & 2 deletions go/vt/vtgate/planbuilder/testdata/aggr_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5690,7 +5690,7 @@
"Variant": "Ordered",
"Aggregates": "group_concat(1) AS group_concat(u.bar), any_value(2) AS baz, any_value(4)",
"GroupBy": "(0|3)",
"ResultColumns": 3,
"ResultColumns": 5,
"Inputs": [
{
"OperatorType": "Join",
Expand Down Expand Up @@ -5957,7 +5957,7 @@
"Variant": "Ordered",
"Aggregates": "sum_count_star(0) AS count(*), any_value(2)",
"GroupBy": "1",
"ResultColumns": 1,
"ResultColumns": 3,
"Inputs": [
{
"OperatorType": "Route",
Expand Down
9 changes: 6 additions & 3 deletions go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,14 @@
"OperatorType": "Sort",
"Variant": "Memory",
"OrderBy": "(1|4) ASC",
"ResultColumns": 3,
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Ordered",
"Aggregates": "any_value(1) AS b, sum_count_star(2) AS count(*), any_value(4)",
"GroupBy": "(0|3)",
"ResultColumns": 3,
"ResultColumns": 5,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -87,13 +88,14 @@
"OperatorType": "Sort",
"Variant": "Memory",
"OrderBy": "(1|4) ASC, (0|3) ASC, 2 ASC",
"ResultColumns": 3,
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Ordered",
"Aggregates": "any_value(1) AS b, sum_count_star(2) AS k, any_value(4)",
"GroupBy": "(0|3)",
"ResultColumns": 3,
"ResultColumns": 5,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -171,13 +173,14 @@
"OperatorType": "Sort",
"Variant": "Memory",
"OrderBy": "(0|3) ASC, 2 ASC",
"ResultColumns": 3,
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Ordered",
"Aggregates": "any_value(1) AS b, sum_count_star(2) AS k",
"GroupBy": "(0|3)",
"ResultColumns": 3,
"ResultColumns": 4,
"Inputs": [
{
"OperatorType": "Route",
Expand Down
13 changes: 8 additions & 5 deletions go/vt/vtgate/planbuilder/testdata/tpch_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,13 +67,14 @@
"OperatorType": "Sort",
"Variant": "Memory",
"OrderBy": "1 DESC COLLATE utf8mb4_0900_ai_ci, (2|5) ASC",
"ResultColumns": 4,
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Ordered",
"Aggregates": "sum(1) AS revenue",
"GroupBy": "(0|4), (2|5), (3|6)",
"ResultColumns": 4,
"ResultColumns": 6,
"Inputs": [
{
"OperatorType": "Projection",
Expand Down Expand Up @@ -1746,7 +1747,7 @@
"Variant": "Ordered",
"Aggregates": "sum_count(1) AS count(o_orderkey), any_value(3)",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"ResultColumns": 4,
"Inputs": [
{
"OperatorType": "Projection",
Expand Down Expand Up @@ -1945,13 +1946,14 @@
"OperatorType": "Sort",
"Variant": "Memory",
"OrderBy": "3 DESC, (0|4) ASC, (1|5) ASC, (2|6) ASC",
"ResultColumns": 4,
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Ordered",
"Aggregates": "count_distinct(3|7) AS supplier_cnt",
"GroupBy": "(0|4), (1|5), (2|6)",
"ResultColumns": 4,
"ResultColumns": 7,
"Inputs": [
{
"OperatorType": "Sort",
Expand Down Expand Up @@ -2084,7 +2086,7 @@
"Variant": "Ordered",
"Aggregates": "sum(5) AS sum(l_quantity)",
"GroupBy": "(4|6), (3|7), (0|8), (1|9), (2|10)",
"ResultColumns": 6,
"ResultColumns": 11,
"Inputs": [
{
"OperatorType": "Sort",
Expand Down Expand Up @@ -2255,13 +2257,14 @@
"OperatorType": "Sort",
"Variant": "Memory",
"OrderBy": "1 DESC, (0|2) ASC",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Aggregate",
"Variant": "Ordered",
"Aggregates": "sum_count_star(1) AS numwait",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"ResultColumns": 3,
"Inputs": [
{
"OperatorType": "Projection",
Expand Down

0 comments on commit 3039c4e

Please sign in to comment.