diff --git a/go/vt/vtgate/planbuilder/operators/aggregator.go b/go/vt/vtgate/planbuilder/operators/aggregator.go index 624a0a297ba..cfeaf1c8cdc 100644 --- a/go/vt/vtgate/planbuilder/operators/aggregator.go +++ b/go/vt/vtgate/planbuilder/operators/aggregator.go @@ -57,6 +57,9 @@ type ( // This is used to truncate the columns in the final result ResultColumns int + // Truncate is set to true if the columns produced by this operator should be truncated if we added any additional columns + Truncate bool + QP *QueryProjection DT *DerivedTable @@ -541,6 +544,12 @@ func (a *Aggregator) pushRemainingGroupingColumnsAndWeightStrings(ctx *planconte } func (a *Aggregator) internalAddWSColumn(ctx *plancontext.PlanningContext, inOffset int, aliasedExpr *sqlparser.AliasedExpr) int { + if a.ResultColumns == 0 && a.Truncate { + // if we need to use `internalAddColumn`, it means we are adding columns that are not part of the original list, + // so we need to set the ResultColumns to the current length of the columns list + a.ResultColumns = len(a.Columns) + } + offset := a.Source.AddWSColumn(ctx, inOffset, false) if offset == len(a.Columns) { @@ -559,7 +568,7 @@ func (a *Aggregator) getTruncateColumnCount() int { } func (a *Aggregator) internalAddColumn(ctx *plancontext.PlanningContext, aliasedExpr *sqlparser.AliasedExpr, addToGroupBy bool) int { - if a.ResultColumns == 0 { + if a.ResultColumns == 0 && a.Truncate { // if we need to use `internalAddColumn`, it means we are adding columns that are not part of the original list, // so we need to set the ResultColumns to the current length of the columns list a.ResultColumns = len(a.Columns) diff --git a/go/vt/vtgate/planbuilder/operators/ast_to_op.go b/go/vt/vtgate/planbuilder/operators/ast_to_op.go index 0d838610866..f017f77d6a3 100644 --- a/go/vt/vtgate/planbuilder/operators/ast_to_op.go +++ b/go/vt/vtgate/planbuilder/operators/ast_to_op.go @@ -148,8 +148,8 @@ func createOperatorFromUnion(ctx *plancontext.PlanningContext, node *sqlparser.U if isRHSUnion { panic(vterrors.VT12001("nesting of UNIONs on the right-hand side")) } - opLHS := translateQueryToOp(ctx, node.Left) - opRHS := translateQueryToOp(ctx, node.Right) + opLHS := translateQueryToOpForUnion(ctx, node.Left) + opRHS := translateQueryToOpForUnion(ctx, node.Right) lexprs := ctx.SemTable.SelectExprs(node.Left) rexprs := ctx.SemTable.SelectExprs(node.Right) @@ -158,6 +158,14 @@ func createOperatorFromUnion(ctx *plancontext.PlanningContext, node *sqlparser.U return newHorizon(union, node) } +func translateQueryToOpForUnion(ctx *plancontext.PlanningContext, node sqlparser.SelectStatement) Operator { + op := translateQueryToOp(ctx, node) + if hz, ok := op.(*Horizon); ok { + hz.Truncate = true + } + return op +} + // createOpFromStmt creates an operator from the given statement. It takes in two additional arguments— // 1. verifyAllFKs: For this given statement, do we need to verify validity of all the foreign keys on the vtgate level. // 2. fkToIgnore: The foreign key constraint to specifically ignore while planning the statement. This field is used in UPDATE CASCADE planning, wherein while planning the child update diff --git a/go/vt/vtgate/planbuilder/operators/horizon.go b/go/vt/vtgate/planbuilder/operators/horizon.go index 7539704c2a9..70186c062b9 100644 --- a/go/vt/vtgate/planbuilder/operators/horizon.go +++ b/go/vt/vtgate/planbuilder/operators/horizon.go @@ -49,6 +49,8 @@ type Horizon struct { // Columns needed to feed other plans Columns []*sqlparser.ColName ColumnsOffset []int + + Truncate bool } func newHorizon(src Operator, query sqlparser.SelectStatement) *Horizon { diff --git a/go/vt/vtgate/planbuilder/operators/horizon_expanding.go b/go/vt/vtgate/planbuilder/operators/horizon_expanding.go index fd3d992d0ca..29c1b1033f1 100644 --- a/go/vt/vtgate/planbuilder/operators/horizon_expanding.go +++ b/go/vt/vtgate/planbuilder/operators/horizon_expanding.go @@ -208,7 +208,7 @@ func createProjectionFromSelect(ctx *plancontext.PlanningContext, horizon *Horiz } if qp.NeedsAggregation() { - return createProjectionWithAggr(ctx, qp, dt, horizon.src()) + return createProjectionWithAggr(ctx, qp, dt, horizon) } projX := createProjectionWithoutAggr(ctx, qp, horizon.src()) @@ -216,8 +216,9 @@ func createProjectionFromSelect(ctx *plancontext.PlanningContext, horizon *Horiz return projX } -func createProjectionWithAggr(ctx *plancontext.PlanningContext, qp *QueryProjection, dt *DerivedTable, src Operator) Operator { +func createProjectionWithAggr(ctx *plancontext.PlanningContext, qp *QueryProjection, dt *DerivedTable, horizon *Horizon) Operator { aggregations, complexAggr := qp.AggregationExpressions(ctx, true) + src := horizon.Source aggrOp := &Aggregator{ Source: src, Original: true, @@ -239,7 +240,11 @@ func createProjectionWithAggr(ctx *plancontext.PlanningContext, qp *QueryProject if complexAggr { return createProjectionForComplexAggregation(aggrOp, qp) } - return createProjectionForSimpleAggregation(ctx, aggrOp, qp) + + addAllColumnsToAggregator(ctx, aggrOp, qp) + aggrOp.Truncate = horizon.Truncate + + return aggrOp } func pullOutValueSubqueries(ctx *plancontext.PlanningContext, aggr Aggr, sqc *SubQueryBuilder, outerID semantics.TableSet) Aggr { @@ -261,7 +266,7 @@ func pullOutValueSubqueries(ctx *plancontext.PlanningContext, aggr Aggr, sqc *Su return aggr } -func createProjectionForSimpleAggregation(ctx *plancontext.PlanningContext, a *Aggregator, qp *QueryProjection) Operator { +func addAllColumnsToAggregator(ctx *plancontext.PlanningContext, a *Aggregator, qp *QueryProjection) { outer: for colIdx, expr := range qp.SelectExprs { ae, err := expr.GetAliasedExpr() @@ -292,7 +297,6 @@ outer: } panic(vterrors.VT13001(fmt.Sprintf("Could not find the %s in aggregation in the original query", sqlparser.String(ae)))) } - return a } func createProjectionForComplexAggregation(a *Aggregator, qp *QueryProjection) Operator { diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json index cd5f052deeb..343159bccaa 100644 --- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json @@ -689,13 +689,13 @@ "OperatorType": "Sort", "Variant": "Memory", "OrderBy": "1 ASC", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "count_distinct(1|3) AS k", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -2149,13 +2149,13 @@ "Instructions": { "OperatorType": "Filter", "Predicate": "count(*) <= 10", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS a", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -2187,13 +2187,13 @@ "Instructions": { "OperatorType": "Filter", "Predicate": "count(*) = 1.00", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum_count_star(0) AS a", "GroupBy": "(1|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -3097,13 +3097,13 @@ "Instructions": { "OperatorType": "Filter", "Predicate": "count(*) = 3", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS count(*)", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -3135,13 +3135,13 @@ "Instructions": { "OperatorType": "Filter", "Predicate": "sum(foo) + sum(bar) = 42", + "ResultColumns": 3, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum(1) AS sum(foo), sum(2) AS sum(bar)", "GroupBy": "(0|3)", - "ResultColumns": 3, "Inputs": [ { "OperatorType": "Route", @@ -3173,13 +3173,13 @@ "Instructions": { "OperatorType": "Filter", "Predicate": "sum(`user`.foo) + sum(bar) = 42", + "ResultColumns": 3, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum(1) AS fooSum, sum(2) AS barSum", "GroupBy": "(0|3)", - "ResultColumns": 3, "Inputs": [ { "OperatorType": "Route", @@ -3218,7 +3218,6 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS count(*)", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -3257,7 +3256,6 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS count(u.`name`)", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -3362,7 +3360,6 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS count(*)", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -4118,7 +4115,6 @@ "Variant": "Ordered", "Aggregates": "max(1|3) AS bazo", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -4157,7 +4153,6 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS bazo", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -5005,13 +5000,13 @@ "Collations": [ "0" ], + "ResultColumns": 1, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum_count_star(0) AS count(*)", "GroupBy": "1", - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -6108,7 +6103,6 @@ "Variant": "Ordered", "Aggregates": "count_star(0)", "GroupBy": "1", - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Projection", @@ -6122,7 +6116,6 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(0) AS count(*)", "GroupBy": "1", - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -6296,7 +6289,6 @@ "Variant": "Ordered", "Aggregates": "sum(1) AS avg(id), sum_count(2) AS count(foo), sum_count(3) AS count(id)", "GroupBy": "(0|4)", - "ResultColumns": 4, "Inputs": [ { "OperatorType": "Route", @@ -6346,7 +6338,6 @@ "Variant": "Ordered", "Aggregates": "sum(1) AS avg(id), sum_count(2) AS count(foo), sum_count(3) AS count(id)", "GroupBy": "(0|4)", - "ResultColumns": 4, "Inputs": [ { "OperatorType": "Route", @@ -6530,7 +6521,6 @@ "OperatorType": "Aggregate", "Variant": "Scalar", "Aggregates": "min(0|2) AS min_id, max(1|3) AS max_id", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json index 82987fb52f4..d6647681103 100644 --- a/go/vt/vtgate/planbuilder/testdata/cte_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/cte_cases.json @@ -485,7 +485,6 @@ "Variant": "Ordered", "Aggregates": "max(1|3) AS bazo", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -524,7 +523,6 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS bazo", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/from_cases.json b/go/vt/vtgate/planbuilder/testdata/from_cases.json index 9a0ffa37f0c..138dc8b045e 100644 --- a/go/vt/vtgate/planbuilder/testdata/from_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/from_cases.json @@ -798,7 +798,6 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS count", "GroupBy": "(0|2)", - "ResultColumns": 3, "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json index 3ce182b51b4..060f073a366 100644 --- a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json @@ -16,7 +16,6 @@ "Variant": "Ordered", "Aggregates": "any_value(1|4) AS b, sum_count_star(2) AS count(*)", "GroupBy": "(0|3)", - "ResultColumns": 5, "Inputs": [ { "OperatorType": "Route", @@ -49,13 +48,13 @@ "OperatorType": "Sort", "Variant": "Memory", "OrderBy": "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, "Inputs": [ { "OperatorType": "Route", @@ -95,7 +94,6 @@ "Variant": "Ordered", "Aggregates": "any_value(1|4) AS b, sum_count_star(2) AS k", "GroupBy": "(0|3)", - "ResultColumns": 5, "Inputs": [ { "OperatorType": "Route", @@ -132,13 +130,13 @@ "OperatorType": "Sort", "Variant": "Memory", "OrderBy": "2 DESC", + "ResultColumns": 3, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "any_value(1) AS b, sum_count_star(2) AS k", "GroupBy": "(0|3)", - "ResultColumns": 3, "Inputs": [ { "OperatorType": "Route", @@ -173,13 +171,13 @@ "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, "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/postprocess_cases.json b/go/vt/vtgate/planbuilder/testdata/postprocess_cases.json index fd8ee10713b..36f1472007d 100644 --- a/go/vt/vtgate/planbuilder/testdata/postprocess_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/postprocess_cases.json @@ -1618,13 +1618,13 @@ "OperatorType": "Sort", "Variant": "Memory", "OrderBy": "0 ASC", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum_count(0) AS count(id)", "GroupBy": "(1|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -2090,7 +2090,6 @@ "Variant": "Ordered", "Aggregates": "min(1|3) AS min(a.id)", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Join", diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json index 1099f62175f..f06a6a50d45 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json @@ -1836,7 +1836,6 @@ "Variant": "Ordered", "Aggregates": "sum(0) AS avg_col, sum_count(3) AS count(intcol)", "GroupBy": "1 COLLATE latin1_swedish_ci, (2|4) COLLATE ", - "ResultColumns": 4, "Inputs": [ { "OperatorType": "Route", @@ -4330,7 +4329,6 @@ "Variant": "Ordered", "Aggregates": "any_value(0) AS id", "GroupBy": "(1|2)", - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -4496,7 +4494,6 @@ "OperatorType": "Aggregate", "Variant": "Scalar", "Aggregates": "max(0|1) AS max(music.id)", - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -5033,7 +5030,6 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS b", "GroupBy": "(2|3), (0|4)", - "ResultColumns": 3, "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/tpch_cases.json b/go/vt/vtgate/planbuilder/testdata/tpch_cases.json index 6e67d4f9e8d..58ec909104a 100644 --- a/go/vt/vtgate/planbuilder/testdata/tpch_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/tpch_cases.json @@ -25,7 +25,6 @@ "Variant": "Ordered", "Aggregates": "sum(2) AS sum_qty, sum(3) AS sum_base_price, sum(4) AS sum_disc_price, sum(5) AS sum_charge, sum(6) AS avg_qty, sum(7) AS avg_price, sum(8) AS avg_disc, sum_count_star(9) AS count_order, sum_count(10) AS count(l_quantity), sum_count(11) AS count(l_extendedprice), sum_count(12) AS count(l_discount)", "GroupBy": "(0|13), (1|14)", - "ResultColumns": 13, "Inputs": [ { "OperatorType": "Route", @@ -67,13 +66,13 @@ "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, "Inputs": [ { "OperatorType": "Projection", @@ -278,13 +277,13 @@ "OperatorType": "Sort", "Variant": "Memory", "OrderBy": "1 DESC COLLATE utf8mb4_0900_ai_ci", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum(1) AS revenue", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -801,7 +800,6 @@ "Variant": "Ordered", "Aggregates": "sum(1) AS sum(case when nation = 'BRAZIL' then volume else 0 end), sum(2) AS sum(volume)", "GroupBy": "(0|3)", - "ResultColumns": 3, "Inputs": [ { "OperatorType": "Projection", @@ -1346,13 +1344,13 @@ "OperatorType": "Sort", "Variant": "Memory", "OrderBy": "2 DESC COLLATE utf8mb4_0900_ai_ci", + "ResultColumns": 8, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum(2) AS revenue", "GroupBy": "(0|8), (1|9), (3|10), (6|11), (4|12), (5|13), (7|14)", - "ResultColumns": 8, "Inputs": [ { "OperatorType": "Projection", @@ -1657,6 +1655,7 @@ "InputName": "Outer", "OperatorType": "Filter", "Predicate": "sum(ps_supplycost * ps_availqty) > :__sq1", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Sort", @@ -1668,7 +1667,6 @@ "Variant": "Ordered", "Aggregates": "sum(1) AS value", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -1900,7 +1898,6 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS count(o_orderkey)", "GroupBy": "(0|2)", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -2061,7 +2058,6 @@ "OperatorType": "Aggregate", "Variant": "Scalar", "Aggregates": "max(0|1) AS max(total_revenue)", - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -2445,13 +2441,13 @@ "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, "Inputs": [ { "OperatorType": "Projection",