Skip to content

Commit

Permalink
limit result columns only when needed
Browse files Browse the repository at this point in the history 
Signed-off-by: Andres Taylor <[email protected]>
  • Loading branch information
@systay
systay committed Aug 14, 2024
1 parent 703dfa4 commit 131b98d
Show file tree
Hide file tree
Showing 11 changed files with 47 additions and 48 deletions.
11 changes: 10 additions & 1 deletion go/vt/vtgate/planbuilder/operators/aggregator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand All @@ -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)
Expand Down
12 changes: 10 additions & 2 deletions go/vt/vtgate/planbuilder/operators/ast_to_op.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go/vt/vtgate/planbuilder/operators/horizon.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
14 changes: 9 additions & 5 deletions go/vt/vtgate/planbuilder/operators/horizon_expanding.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -208,16 +208,17 @@ 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())
projX.DT = dt
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,
Expand All @@ -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 {
Expand All @@ -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()
Expand Down Expand Up @@ -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 {
Expand Down
24 changes: 7 additions & 17 deletions go/vt/vtgate/planbuilder/testdata/aggr_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -3218,7 +3218,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count_star(1) AS count(*)",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -3257,7 +3256,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count(1) AS count(u.`name`)",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Projection",
Expand Down Expand Up @@ -3362,7 +3360,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count_star(1) AS count(*)",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Projection",
Expand Down Expand Up @@ -4118,7 +4115,6 @@
"Variant": "Ordered",
"Aggregates": "max(1|3) AS bazo",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -4157,7 +4153,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count(1) AS bazo",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -6108,7 +6103,6 @@
"Variant": "Ordered",
"Aggregates": "count_star(0)",
"GroupBy": "1",
"ResultColumns": 1,
"Inputs": [
{
"OperatorType": "Projection",
Expand All @@ -6122,7 +6116,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count_star(0) AS count(*)",
"GroupBy": "1",
"ResultColumns": 1,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
2 changes: 0 additions & 2 deletions go/vt/vtgate/planbuilder/testdata/cte_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -485,7 +485,6 @@
"Variant": "Ordered",
"Aggregates": "max(1|3) AS bazo",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -524,7 +523,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count(1) AS bazo",
"GroupBy": "(0|2)",
"ResultColumns": 2,
"Inputs": [
{
"OperatorType": "Route",
Expand Down
1 change: 0 additions & 1 deletion go/vt/vtgate/planbuilder/testdata/from_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -798,7 +798,6 @@
"Variant": "Ordered",
"Aggregates": "sum_count_star(1) AS count",
"GroupBy": "(0|2)",
"ResultColumns": 3,
"Inputs": [
{
"OperatorType": "Route",
Expand Down
8 changes: 3 additions & 5 deletions go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
Loading

0 comments on commit 131b98d

Please sign in to comment.