From 6441298015182960d9569e09d0f35ccb1a5bfd8b Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Wed, 19 Jun 2024 08:46:59 +0200 Subject: [PATCH 1/3] bugfix: handle literals in UNIONs better Signed-off-by: Andres Taylor --- .../planbuilder/operator_transformers.go | 4 +- .../planbuilder/operators/aggregator.go | 23 +++++++-- .../vtgate/planbuilder/operators/distinct.go | 14 ++++-- go/vt/vtgate/planbuilder/operators/filter.go | 14 ++++-- .../vtgate/planbuilder/operators/ordering.go | 8 ++- .../planbuilder/operators/plan_query.go | 24 +++++++-- go/vt/vtgate/planbuilder/operators/route.go | 8 ++- go/vt/vtgate/planbuilder/operators/union.go | 5 +- .../planbuilder/testdata/aggr_cases.json | 27 +++++++--- .../planbuilder/testdata/cte_cases.json | 3 ++ .../testdata/memory_sort_cases.json | 10 ++-- .../testdata/postprocess_cases.json | 3 +- .../planbuilder/testdata/select_cases.json | 4 ++ .../planbuilder/testdata/tpch_cases.json | 17 ++++--- .../planbuilder/testdata/union_cases.json | 49 +++++++++++++++++++ 15 files changed, 168 insertions(+), 45 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go index bec5cd28bb5..60a59ed936a 100644 --- a/go/vt/vtgate/planbuilder/operator_transformers.go +++ b/go/vt/vtgate/planbuilder/operator_transformers.go @@ -352,7 +352,7 @@ func transformDistinct(ctx *plancontext.PlanningContext, op *operators.Distinct) return &engine.Distinct{ Source: src, CheckCols: op.Columns, - Truncate: op.Truncate, + Truncate: op.ResultColumns, }, nil } @@ -470,7 +470,7 @@ func transformFilter(ctx *plancontext.PlanningContext, op *operators.Filter) (en Input: src, Predicate: predicate, ASTPredicate: ctx.SemTable.AndExpressions(op.Predicates...), - Truncate: op.Truncate, + Truncate: op.ResultColumns, }, nil } diff --git a/go/vt/vtgate/planbuilder/operators/aggregator.go b/go/vt/vtgate/planbuilder/operators/aggregator.go index e9fee905024..e198b4417a5 100644 --- a/go/vt/vtgate/planbuilder/operators/aggregator.go +++ b/go/vt/vtgate/planbuilder/operators/aggregator.go @@ -51,7 +51,10 @@ type ( offsetPlanned bool // Original will only be true for the original aggregator created from the AST - Original bool + Original bool + + // ResultColumns signals how many columns will be produced by this operator + // This is used to truncate the columns in the final result ResultColumns int QP *QueryProjection @@ -281,9 +284,9 @@ func isDerived(op Operator) bool { } } -func (a *Aggregator) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr { +func (a *Aggregator) GetColumns(ctx *plancontext.PlanningContext) (res []*sqlparser.AliasedExpr) { if isDerived(a.Source) { - return a.Columns + return truncate(a, a.Columns) } // we update the incoming columns, so we know about any new columns that have been added @@ -296,7 +299,7 @@ func (a *Aggregator) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.A a.Columns = append(a.Columns, columns[len(a.Columns):]...) } - return a.Columns + return truncate(a, a.Columns) } func (a *Aggregator) GetSelectExprs(ctx *plancontext.PlanningContext) sqlparser.SelectExprs { @@ -315,6 +318,9 @@ func (a *Aggregator) ShortDescription() string { if a.Original { org = "ORG " } + if a.ResultColumns > 0 { + org += fmt.Sprintf(":%d ", a.ResultColumns) + } if len(a.Grouping) == 0 { return fmt.Sprintf("%s%s", org, strings.Join(columns, ", ")) @@ -511,7 +517,16 @@ func (a *Aggregator) setTruncateColumnCount(offset int) { a.ResultColumns = offset } +func (a *Aggregator) getTruncateColumnCount() int { + return a.ResultColumns +} + func (a *Aggregator) internalAddColumn(ctx *plancontext.PlanningContext, aliasedExpr *sqlparser.AliasedExpr, addToGroupBy bool) int { + if a.ResultColumns == 0 { + // 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.AddColumn(ctx, true, addToGroupBy, aliasedExpr) if offset == len(a.Columns) { diff --git a/go/vt/vtgate/planbuilder/operators/distinct.go b/go/vt/vtgate/planbuilder/operators/distinct.go index f24d5b4978b..7807b94d491 100644 --- a/go/vt/vtgate/planbuilder/operators/distinct.go +++ b/go/vt/vtgate/planbuilder/operators/distinct.go @@ -41,7 +41,7 @@ type ( // This is only filled in during offset planning Columns []engine.CheckCol - Truncate int + ResultColumns int } ) @@ -72,7 +72,7 @@ func (d *Distinct) Clone(inputs []Operator) Operator { Columns: slices.Clone(d.Columns), QP: d.QP, PushedPerformance: d.PushedPerformance, - Truncate: d.Truncate, + ResultColumns: d.ResultColumns, } } @@ -101,11 +101,11 @@ func (d *Distinct) FindCol(ctx *plancontext.PlanningContext, expr sqlparser.Expr } func (d *Distinct) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr { - return d.Source.GetColumns(ctx) + return truncate(d, d.Source.GetColumns(ctx)) } func (d *Distinct) GetSelectExprs(ctx *plancontext.PlanningContext) sqlparser.SelectExprs { - return d.Source.GetSelectExprs(ctx) + return truncate(d, d.Source.GetSelectExprs(ctx)) } func (d *Distinct) ShortDescription() string { @@ -120,5 +120,9 @@ func (d *Distinct) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy { } func (d *Distinct) setTruncateColumnCount(offset int) { - d.Truncate = offset + d.ResultColumns = offset +} + +func (d *Distinct) getTruncateColumnCount() int { + return d.ResultColumns } diff --git a/go/vt/vtgate/planbuilder/operators/filter.go b/go/vt/vtgate/planbuilder/operators/filter.go index 19d864c0ada..d68b2a43a24 100644 --- a/go/vt/vtgate/planbuilder/operators/filter.go +++ b/go/vt/vtgate/planbuilder/operators/filter.go @@ -36,7 +36,7 @@ type Filter struct { // It contains the ANDed predicates in Predicates, with ColName:s replaced by Offset:s PredicateWithOffsets evalengine.Expr - Truncate int + ResultColumns int } func newFilterSinglePredicate(op Operator, expr sqlparser.Expr) Operator { @@ -55,7 +55,7 @@ func (f *Filter) Clone(inputs []Operator) Operator { Source: inputs[0], Predicates: slices.Clone(f.Predicates), PredicateWithOffsets: f.PredicateWithOffsets, - Truncate: f.Truncate, + ResultColumns: f.ResultColumns, } } @@ -100,11 +100,11 @@ func (f *Filter) AddWSColumn(ctx *plancontext.PlanningContext, offset int, under } func (f *Filter) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr { - return f.Source.GetColumns(ctx) + return truncate(f, f.Source.GetColumns(ctx)) } func (f *Filter) GetSelectExprs(ctx *plancontext.PlanningContext) sqlparser.SelectExprs { - return f.Source.GetSelectExprs(ctx) + return truncate(f, f.Source.GetSelectExprs(ctx)) } func (f *Filter) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy { @@ -151,5 +151,9 @@ func (f *Filter) ShortDescription() string { } func (f *Filter) setTruncateColumnCount(offset int) { - f.Truncate = offset + f.ResultColumns = offset +} + +func (f *Filter) getTruncateColumnCount() int { + return f.ResultColumns } diff --git a/go/vt/vtgate/planbuilder/operators/ordering.go b/go/vt/vtgate/planbuilder/operators/ordering.go index f8008022511..5414b34fc40 100644 --- a/go/vt/vtgate/planbuilder/operators/ordering.go +++ b/go/vt/vtgate/planbuilder/operators/ordering.go @@ -70,11 +70,11 @@ func (o *Ordering) FindCol(ctx *plancontext.PlanningContext, expr sqlparser.Expr } func (o *Ordering) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr { - return o.Source.GetColumns(ctx) + return truncate(o, o.Source.GetColumns(ctx)) } func (o *Ordering) GetSelectExprs(ctx *plancontext.PlanningContext) sqlparser.SelectExprs { - return o.Source.GetSelectExprs(ctx) + return truncate(o, o.Source.GetSelectExprs(ctx)) } func (o *Ordering) GetOrdering(*plancontext.PlanningContext) []OrderBy { @@ -108,3 +108,7 @@ func (o *Ordering) ShortDescription() string { func (o *Ordering) setTruncateColumnCount(offset int) { o.ResultColumns = offset } + +func (o *Ordering) getTruncateColumnCount() int { + return o.ResultColumns +} diff --git a/go/vt/vtgate/planbuilder/operators/plan_query.go b/go/vt/vtgate/planbuilder/operators/plan_query.go index ea6b88f752d..4d371942c26 100644 --- a/go/vt/vtgate/planbuilder/operators/plan_query.go +++ b/go/vt/vtgate/planbuilder/operators/plan_query.go @@ -131,12 +131,21 @@ func (noPredicates) AddPredicate(*plancontext.PlanningContext, sqlparser.Expr) O panic(vterrors.VT13001("the noColumns operator cannot accept predicates")) } -// tryTruncateColumnsAt will see if we can truncate the columns by just asking the operator to do it for us -func tryTruncateColumnsAt(op Operator, truncateAt int) bool { - type columnTruncator interface { - setTruncateColumnCount(offset int) +// columnTruncator is an interface that allows an operator to truncate its columns to a certain length +type columnTruncator interface { + setTruncateColumnCount(offset int) + getTruncateColumnCount() int +} + +func truncate[K any](op columnTruncator, slice []K) []K { + if op.getTruncateColumnCount() == 0 { + return slice } + return slice[:op.getTruncateColumnCount()] +} +// tryTruncateColumnsAt will see if we can truncate the columns by just asking the operator to do it for us +func tryTruncateColumnsAt(op Operator, truncateAt int) bool { truncator, ok := op.(columnTruncator) if ok { truncator.setTruncateColumnCount(truncateAt) @@ -160,6 +169,13 @@ func tryTruncateColumnsAt(op Operator, truncateAt int) bool { func transformColumnsToSelectExprs(ctx *plancontext.PlanningContext, op Operator) sqlparser.SelectExprs { columns := op.GetColumns(ctx) + if trunc, ok := op.(columnTruncator); ok { + count := trunc.getTruncateColumnCount() + if count > 0 { + columns = columns[:count] + } + } + selExprs := slice.Map(columns, func(from *sqlparser.AliasedExpr) sqlparser.SelectExpr { return from }) diff --git a/go/vt/vtgate/planbuilder/operators/route.go b/go/vt/vtgate/planbuilder/operators/route.go index feeb091a725..cc049d22753 100644 --- a/go/vt/vtgate/planbuilder/operators/route.go +++ b/go/vt/vtgate/planbuilder/operators/route.go @@ -749,11 +749,11 @@ func (r *Route) FindCol(ctx *plancontext.PlanningContext, expr sqlparser.Expr, _ } func (r *Route) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr { - return r.Source.GetColumns(ctx) + return truncate(r, r.Source.GetColumns(ctx)) } func (r *Route) GetSelectExprs(ctx *plancontext.PlanningContext) sqlparser.SelectExprs { - return r.Source.GetSelectExprs(ctx) + return truncate(r, r.Source.GetSelectExprs(ctx)) } func (r *Route) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy { @@ -849,6 +849,10 @@ func (r *Route) setTruncateColumnCount(offset int) { r.ResultColumns = offset } +func (r *Route) getTruncateColumnCount() int { + return r.ResultColumns +} + func (r *Route) introducesTableID() semantics.TableSet { id := semantics.EmptyTableSet() for _, route := range r.MergedWith { diff --git a/go/vt/vtgate/planbuilder/operators/union.go b/go/vt/vtgate/planbuilder/operators/union.go index fedfc362017..7d09391cf7d 100644 --- a/go/vt/vtgate/planbuilder/operators/union.go +++ b/go/vt/vtgate/planbuilder/operators/union.go @@ -208,7 +208,8 @@ func (u *Union) addConstantToUnion(ctx *plancontext.PlanningContext, aexpr *sqlp outputOffset = thisOffset } else { if thisOffset != outputOffset { - panic(vterrors.VT12001("argument offsets did not line up for UNION")) + tree := ToTree(u) + panic(vterrors.VT12001(fmt.Sprintf("argument offsets did not line up for UNION. Pushing %s - want %d got %d\n%s", sqlparser.String(aexpr), outputOffset, thisOffset, tree))) } } } @@ -224,7 +225,7 @@ func (u *Union) addWeightStringToOffset(ctx *plancontext.PlanningContext, argIdx outputOffset = thisOffset } else { if thisOffset != outputOffset { - panic(vterrors.VT12001("weight_string offsets did not line up for UNION")) + panic(vterrors.VT13001("weight_string offsets did not line up for UNION")) } } } diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json index fb04edd6d44..f1017a51e4a 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", @@ -2039,13 +2039,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", @@ -2077,13 +2077,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", @@ -2939,13 +2939,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", @@ -2977,13 +2977,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", @@ -3015,13 +3015,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", @@ -3060,6 +3060,7 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS count(*)", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -3098,6 +3099,7 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS count(u.`name`)", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -3202,6 +3204,7 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(1) AS count(*)", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -3951,6 +3954,7 @@ "Variant": "Ordered", "Aggregates": "max(1|3) AS bazo", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -3989,6 +3993,7 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS bazo", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -4836,13 +4841,13 @@ "Collations": [ "0" ], - "ResultColumns": 1, "Inputs": [ { "OperatorType": "Aggregate", "Variant": "Ordered", "Aggregates": "sum_count_star(0) AS count(*)", "GroupBy": "1", + "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -5518,6 +5523,7 @@ "OperatorType": "Aggregate", "Variant": "Ordered", "GroupBy": "0, (1|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "SimpleProjection", @@ -5684,6 +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, "Inputs": [ { "OperatorType": "Join", @@ -5939,6 +5946,7 @@ "Variant": "Ordered", "Aggregates": "count_star(0)", "GroupBy": "1", + "ResultColumns": 1, "Inputs": [ { "OperatorType": "SimpleProjection", @@ -5949,6 +5957,7 @@ "Variant": "Ordered", "Aggregates": "sum_count_star(0) AS count(*), any_value(2)", "GroupBy": "1", + "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -6122,6 +6131,7 @@ "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", @@ -6171,6 +6181,7 @@ "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", @@ -6354,6 +6365,7 @@ "OperatorType": "Aggregate", "Variant": "Scalar", "Aggregates": "min(0|2) AS min_id, max(1|2) AS max_id", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -6517,6 +6529,7 @@ "Variant": "Ordered", "Aggregates": "count_star(0)", "GroupBy": "1, (2|3)", + "ResultColumns": 3, "Inputs": [ { "OperatorType": "SimpleProjection", diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json index 8181c13cd0b..09d155b19f6 100644 --- a/go/vt/vtgate/planbuilder/testdata/cte_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/cte_cases.json @@ -479,6 +479,7 @@ "Variant": "Ordered", "Aggregates": "max(1|3) AS bazo", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -517,6 +518,7 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS bazo", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Route", @@ -563,6 +565,7 @@ "OperatorType": "Aggregate", "Variant": "Ordered", "GroupBy": "0, (1|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "SimpleProjection", diff --git a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json index 34f198abb96..e17ceb3524c 100644 --- a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json @@ -9,13 +9,13 @@ "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, "Inputs": [ { "OperatorType": "Route", @@ -48,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", @@ -87,13 +87,13 @@ "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, "Inputs": [ { "OperatorType": "Route", @@ -130,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", @@ -171,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 96a92d5894d..454740f0498 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,6 +2090,7 @@ "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 6dfa03b79e7..38f0cb3044f 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json @@ -1752,6 +1752,7 @@ "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", @@ -4065,6 +4066,7 @@ "Variant": "Ordered", "Aggregates": "any_value(0) AS id", "GroupBy": "(1|2)", + "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -4230,6 +4232,7 @@ "OperatorType": "Aggregate", "Variant": "Scalar", "Aggregates": "max(0|1) AS max(music.id)", + "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -4766,6 +4769,7 @@ "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 4a9990b0e9b..6bc14f1625b 100644 --- a/go/vt/vtgate/planbuilder/testdata/tpch_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/tpch_cases.json @@ -25,6 +25,7 @@ "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", @@ -66,13 +67,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", @@ -277,13 +278,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", @@ -741,6 +742,7 @@ "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": "SimpleProjection", @@ -1204,13 +1206,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", @@ -1515,7 +1517,6 @@ "InputName": "Outer", "OperatorType": "Filter", "Predicate": "sum(ps_supplycost * ps_availqty) > :__sq1", - "ResultColumns": 2, "Inputs": [ { "OperatorType": "Sort", @@ -1527,6 +1528,7 @@ "Variant": "Ordered", "Aggregates": "sum(1) AS value", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -1744,6 +1746,7 @@ "Variant": "Ordered", "Aggregates": "sum_count(1) AS count(o_orderkey), any_value(3)", "GroupBy": "(0|2)", + "ResultColumns": 2, "Inputs": [ { "OperatorType": "Projection", @@ -1895,6 +1898,7 @@ "OperatorType": "Aggregate", "Variant": "Scalar", "Aggregates": "max(0|1) AS max(total_revenue)", + "ResultColumns": 1, "Inputs": [ { "OperatorType": "Route", @@ -1941,13 +1945,13 @@ "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, "Inputs": [ { "OperatorType": "Sort", @@ -2080,6 +2084,7 @@ "Variant": "Ordered", "Aggregates": "sum(5) AS sum(l_quantity)", "GroupBy": "(4|6), (3|7), (0|8), (1|9), (2|10)", + "ResultColumns": 6, "Inputs": [ { "OperatorType": "Sort", @@ -2250,13 +2255,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", diff --git a/go/vt/vtgate/planbuilder/testdata/union_cases.json b/go/vt/vtgate/planbuilder/testdata/union_cases.json index 3cd698342a5..49458f8c608 100644 --- a/go/vt/vtgate/planbuilder/testdata/union_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/union_cases.json @@ -1808,5 +1808,54 @@ "user.user" ] } + }, + { + "comment": "Literals in UNION can be problematic", + "query": "select 'a' as type, 0 as id from user group by 2 union all select 'c' as type, 0 as id from user_extra as t", + "plan": { + "QueryType": "SELECT", + "Original": "select 'a' as type, 0 as id from user group by 2 union all select 'c' as type, 0 as id from user_extra as t", + "Instructions": { + "OperatorType": "Concatenate", + "Inputs": [ + { + "OperatorType": "Aggregate", + "Variant": "Ordered", + "Aggregates": "any_value(0) AS type, any_value(1) AS id", + "GroupBy": "2 COLLATE utf8mb4_0900_ai_ci", + "ResultColumns": 2, + "Inputs": [ + { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select 'a' as type, 0 as id, '' from `user` where 1 != 1 group by ''", + "OrderBy": "2 ASC COLLATE utf8mb4_0900_ai_ci", + "Query": "select 'a' as type, 0 as id, '' from `user` group by '' order by '' asc", + "Table": "`user`" + } + ] + }, + { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select 'c' as type, 0 as id from user_extra as t where 1 != 1", + "Query": "select 'c' as type, 0 as id from user_extra as t", + "Table": "user_extra" + } + ] + }, + "TablesUsed": [ + "user.user", + "user.user_extra" + ] + } } ] From 3039c4e0c82b47bc0f1724b13040a9ae56d752a4 Mon Sep 17 00:00:00 2001 From: Florent Poinsard Date: Wed, 19 Jun 2024 09:00:23 -0600 Subject: [PATCH 2/3] Stop truncating needed offsets Once an offset has been shared with the outside, we need to not truncate away that column Signed-off-by: Florent Poinsard Signed-off-by: Andres Taylor --- go/vt/vtgate/planbuilder/operators/aggregator.go | 10 ++++++++++ go/vt/vtgate/planbuilder/testdata/aggr_cases.json | 4 ++-- .../planbuilder/testdata/memory_sort_cases.json | 9 ++++++--- go/vt/vtgate/planbuilder/testdata/tpch_cases.json | 13 ++++++++----- 4 files changed, 26 insertions(+), 10 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operators/aggregator.go b/go/vt/vtgate/planbuilder/operators/aggregator.go index e198b4417a5..73ec9eb5474 100644 --- a/go/vt/vtgate/planbuilder/operators/aggregator.go +++ b/go/vt/vtgate/planbuilder/operators/aggregator.go @@ -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) @@ -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 } } @@ -191,6 +199,7 @@ func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gro panic(errFailedToPlan(ae)) } + a.checkOffset(offset) return offset } @@ -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 } diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json index f1017a51e4a..dac95cfa51f 100644 --- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json @@ -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", @@ -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", diff --git a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json index e17ceb3524c..3281feacc76 100644 --- a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json @@ -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", @@ -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", @@ -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", diff --git a/go/vt/vtgate/planbuilder/testdata/tpch_cases.json b/go/vt/vtgate/planbuilder/testdata/tpch_cases.json index 6bc14f1625b..610a0ba1c23 100644 --- a/go/vt/vtgate/planbuilder/testdata/tpch_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/tpch_cases.json @@ -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", @@ -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", @@ -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", @@ -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", @@ -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", From 826abd65f1b267cf504607794c8931d100f6297b Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Thu, 20 Jun 2024 07:41:12 +0200 Subject: [PATCH 3/3] review touch ups Signed-off-by: Andres Taylor --- go/vt/vtgate/planbuilder/operators/aggregator.go | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operators/aggregator.go b/go/vt/vtgate/planbuilder/operators/aggregator.go index 73ec9eb5474..40ba20fe02d 100644 --- a/go/vt/vtgate/planbuilder/operators/aggregator.go +++ b/go/vt/vtgate/planbuilder/operators/aggregator.go @@ -151,12 +151,17 @@ func (a *Aggregator) FindCol(ctx *plancontext.PlanningContext, in sqlparser.Expr } func (a *Aggregator) checkOffset(offset int) { + // if the offset is greater than the number of columns we expect to produce, we need to update the number of columns + // this is to make sure that the column is not truncated in the final result 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 { +func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, groupBy bool, ae *sqlparser.AliasedExpr) (offset int) { + defer func() { + a.checkOffset(offset) + }() rewritten := a.DT.RewriteExpression(ctx, ae.Expr) ae = &sqlparser.AliasedExpr{ @@ -167,7 +172,6 @@ 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 } } @@ -191,7 +195,7 @@ func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gro a.Aggregations = append(a.Aggregations, aggr) } - offset := len(a.Columns) + offset = len(a.Columns) a.Columns = append(a.Columns, ae) incomingOffset := a.Source.AddColumn(ctx, false, groupBy, ae) @@ -199,7 +203,6 @@ func (a *Aggregator) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gro panic(errFailedToPlan(ae)) } - a.checkOffset(offset) return offset }