vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/memory_sort.go (about)

     1  /*
     2  Copyright 2019 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/vtgate/planbuilder/plancontext"
    23  
    24  	"vitess.io/vitess/go/mysql/collations"
    25  	"vitess.io/vitess/go/vt/vterrors"
    26  
    27  	"vitess.io/vitess/go/sqltypes"
    28  	"vitess.io/vitess/go/vt/sqlparser"
    29  	"vitess.io/vitess/go/vt/vtgate/engine"
    30  )
    31  
    32  var _ logicalPlan = (*memorySort)(nil)
    33  
    34  // memorySort is the logicalPlan for engine.Limit.
    35  // This gets built if a limit needs to be applied
    36  // after rows are returned from an underlying
    37  // operation. Since a limit is the final operation
    38  // of a SELECT, most pushes are not applicable.
    39  type memorySort struct {
    40  	resultsBuilder
    41  	eMemorySort *engine.MemorySort
    42  }
    43  
    44  func findColNumber(ms *memorySort, expr *sqlparser.ColName) int {
    45  	c := expr.Metadata.(*column)
    46  	for i, rc := range ms.ResultColumns() {
    47  		if rc.column == c {
    48  			return i
    49  		}
    50  	}
    51  	return -1
    52  }
    53  
    54  // newMemorySort builds a new memorySort.
    55  func newMemorySort(plan logicalPlan, orderBy v3OrderBy) (*memorySort, error) {
    56  	eMemorySort := &engine.MemorySort{}
    57  	ms := &memorySort{
    58  		resultsBuilder: newResultsBuilder(plan, eMemorySort),
    59  		eMemorySort:    eMemorySort,
    60  	}
    61  	for _, order := range orderBy {
    62  		var colNumber int
    63  		switch expr := order.Expr.(type) {
    64  		case *sqlparser.Literal:
    65  			var err error
    66  			if colNumber, err = ResultFromNumber(ms.ResultColumns(), expr, "order clause"); err != nil {
    67  				return nil, err
    68  			}
    69  		case *sqlparser.ColName:
    70  			colNumber = findColNumber(ms, expr)
    71  		case *sqlparser.CastExpr:
    72  			colName, ok := expr.Expr.(*sqlparser.ColName)
    73  			if !ok {
    74  				return nil, vterrors.VT12001(fmt.Sprintf("memory sort: complex ORDER BY expression: %s", sqlparser.String(expr)))
    75  			}
    76  			colNumber = findColNumber(ms, colName)
    77  		case *sqlparser.ConvertExpr:
    78  			colName, ok := expr.Expr.(*sqlparser.ColName)
    79  			if !ok {
    80  				return nil, vterrors.VT12001(fmt.Sprintf("memory sort: complex ORDER BY expression: %s", sqlparser.String(expr)))
    81  			}
    82  			colNumber = findColNumber(ms, colName)
    83  		default:
    84  			return nil, vterrors.VT12001(fmt.Sprintf("memory sort: complex ORDER BY expression: %s", sqlparser.String(expr)))
    85  		}
    86  		// If column is not found, then the order by is referencing
    87  		// a column that's not on the select list.
    88  		if colNumber == -1 {
    89  			return nil, vterrors.VT12001(fmt.Sprintf("memory sort: ORDER BY must reference a column in the SELECT list: %s", sqlparser.String(order)))
    90  		}
    91  		// TODO(king-11) need to pass in collation here
    92  		ob := engine.OrderByParams{
    93  			Col:               colNumber,
    94  			WeightStringCol:   -1,
    95  			Desc:              order.Direction == sqlparser.DescOrder,
    96  			StarColFixedIndex: colNumber,
    97  			FromGroupBy:       order.fromGroupBy,
    98  			CollationID:       collations.Unknown,
    99  		}
   100  		ms.eMemorySort.OrderBy = append(ms.eMemorySort.OrderBy, ob)
   101  	}
   102  	return ms, nil
   103  }
   104  
   105  // Primitive implements the logicalPlan interface
   106  func (ms *memorySort) Primitive() engine.Primitive {
   107  	ms.eMemorySort.Input = ms.input.Primitive()
   108  	return ms.eMemorySort
   109  }
   110  
   111  // SetLimit implements the logicalPlan interface
   112  func (ms *memorySort) SetLimit(limit *sqlparser.Limit) error {
   113  	return vterrors.VT13001("memorySort.Limit: unreachable")
   114  }
   115  
   116  // Wireup implements the logicalPlan interface
   117  // If text columns are detected in the keys, then the function modifies
   118  // the primitive to pull a corresponding weight_string from mysql and
   119  // compare those instead. This is because we currently don't have the
   120  // ability to mimic mysql's collation behavior.
   121  func (ms *memorySort) Wireup(plan logicalPlan, jt *jointab) error {
   122  	for i, orderby := range ms.eMemorySort.OrderBy {
   123  		rc := ms.resultColumns[orderby.Col]
   124  		// Add a weight_string column if we know that the column is a textual column or if its type is unknown
   125  		if sqltypes.IsText(rc.column.typ) || rc.column.typ == sqltypes.Null {
   126  			weightcolNumber, err := ms.input.SupplyWeightString(orderby.Col, orderby.FromGroupBy)
   127  			if err != nil {
   128  				_, isUnsupportedErr := err.(UnsupportedSupplyWeightString)
   129  				if isUnsupportedErr {
   130  					continue
   131  				}
   132  				return err
   133  			}
   134  			ms.weightStrings[rc] = weightcolNumber
   135  			ms.eMemorySort.OrderBy[i].WeightStringCol = weightcolNumber
   136  			ms.eMemorySort.TruncateColumnCount = len(ms.resultColumns)
   137  		}
   138  	}
   139  	return ms.input.Wireup(plan, jt)
   140  }
   141  
   142  func (ms *memorySort) WireupGen4(ctx *plancontext.PlanningContext) error {
   143  	return ms.input.WireupGen4(ctx)
   144  }