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  }