vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/projection.go (about) 1 /* 2 Copyright 2022 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package planbuilder 18 19 import ( 20 "fmt" 21 22 "vitess.io/vitess/go/vt/sqlparser" 23 "vitess.io/vitess/go/vt/vterrors" 24 "vitess.io/vitess/go/vt/vtgate/engine" 25 "vitess.io/vitess/go/vt/vtgate/evalengine" 26 "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" 27 "vitess.io/vitess/go/vt/vtgate/semantics" 28 ) 29 30 type projection struct { 31 gen4Plan 32 source logicalPlan 33 columnNames []string 34 columns []sqlparser.Expr 35 primitive *engine.Projection 36 // unorderedColumnIdx is used to find the index at which we should add any column output from projection 37 // we don't care for the ordering of. It should also be updated when such a column is added 38 unorderedColumnIdx int 39 } 40 41 var _ logicalPlan = (*projection)(nil) 42 43 // WireupGen4 implements the logicalPlan interface 44 func (p *projection) WireupGen4(ctx *plancontext.PlanningContext) error { 45 columns := make([]evalengine.Expr, 0, len(p.columns)) 46 for _, expr := range p.columns { 47 convert, err := evalengine.Translate(expr, ctx.SemTable) 48 if err != nil { 49 return err 50 } 51 columns = append(columns, convert) 52 } 53 p.primitive = &engine.Projection{ 54 Cols: p.columnNames, 55 Exprs: columns, 56 } 57 58 return p.source.WireupGen4(ctx) 59 } 60 61 // Inputs implements the logicalPlan interface 62 func (p *projection) Inputs() []logicalPlan { 63 return []logicalPlan{p.source} 64 } 65 66 // Rewrite implements the logicalPlan interface 67 func (p *projection) Rewrite(inputs ...logicalPlan) error { 68 if len(inputs) != 1 { 69 return vterrors.VT13001(fmt.Sprintf("wrong number of inputs, got: %d; expected: %d", len(inputs), 1)) 70 } 71 p.source = inputs[0] 72 return nil 73 } 74 75 // ContainsTables implements the logicalPlan interface 76 func (p *projection) ContainsTables() semantics.TableSet { 77 return p.source.ContainsTables() 78 } 79 80 // OutputColumns implements the logicalPlan interface 81 func (p *projection) OutputColumns() []sqlparser.SelectExpr { 82 columns := make([]sqlparser.SelectExpr, 0, len(p.columns)) 83 for i, expr := range p.columns { 84 columns = append(columns, &sqlparser.AliasedExpr{ 85 Expr: expr, 86 As: sqlparser.NewIdentifierCI(p.columnNames[i]), 87 }) 88 } 89 return columns 90 } 91 92 // Primitive implements the logicalPlan interface 93 func (p *projection) Primitive() engine.Primitive { 94 if p.primitive == nil { 95 panic("WireUp not yet run") 96 } 97 p.primitive.Input = p.source.Primitive() 98 return p.primitive 99 } 100 101 // addColumn is used to add a column output for the projection. 102 // This is the only function that should be used to add columns to projection 103 func (p *projection) addColumn(idx *int, column sqlparser.Expr, columnName string) (int, error) { 104 var offset int 105 if idx == nil { 106 p.unorderedColumnIdx++ 107 offset = len(p.columns) - p.unorderedColumnIdx 108 } else { 109 offset = *idx 110 } 111 if p.columnNames[offset] != "" || p.columns[offset] != nil { 112 return -1, vterrors.VT13001("overwriting columns in projection is not permitted") 113 } 114 p.columns[offset] = column 115 p.columnNames[offset] = columnName 116 return offset, nil 117 }