vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/operators/union.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 operators
    18  
    19  import (
    20  	"vitess.io/vitess/go/vt/sqlparser"
    21  	"vitess.io/vitess/go/vt/vterrors"
    22  	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/ops"
    23  	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/rewrite"
    24  	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
    25  )
    26  
    27  type Union struct {
    28  	Sources  []ops.Operator
    29  	Distinct bool
    30  
    31  	// TODO this should be removed. For now it's used to fail queries
    32  	Ordering sqlparser.OrderBy
    33  
    34  	noColumns
    35  }
    36  
    37  var _ ops.PhysicalOperator = (*Union)(nil)
    38  
    39  // IPhysical implements the PhysicalOperator interface
    40  func (u *Union) IPhysical() {}
    41  
    42  // Clone implements the Operator interface
    43  func (u *Union) Clone(inputs []ops.Operator) ops.Operator {
    44  	newOp := *u
    45  	newOp.Sources = inputs
    46  	return &newOp
    47  }
    48  
    49  // Inputs implements the Operator interface
    50  func (u *Union) Inputs() []ops.Operator {
    51  	return u.Sources
    52  }
    53  
    54  // AddPredicate adds a predicate a UNION by pushing the predicate to all sources of the UNION.
    55  /* this is done by offset and expression rewriting. Say we have a query like so:
    56  select * (
    57   	select foo as col, bar from tbl1
    58  	union
    59  	select id, baz from tbl2
    60  ) as X where X.col = 42
    61  
    62  We want to push down the `X.col = 42` as far down the operator tree as possible. We want
    63  to end up with an operator tree that looks something like this:
    64  
    65  select * (
    66   	select foo as col, bar from tbl1 where foo = 42
    67  	union
    68  	select id, baz from tbl2 where id = 42
    69  ) as X
    70  
    71  Notice how `X.col = 42` has been translated to `foo = 42` and `id = 42` on respective WHERE clause.
    72  The first SELECT of the union dictates the column names, and the second is whatever expression
    73  can be found on the same offset. The names of the RHS are discarded.
    74  */
    75  func (u *Union) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Expr) (ops.Operator, error) {
    76  	offsets := make(map[string]int)
    77  	sel, err := u.GetSelectFor(0)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	for i, selectExpr := range sel.SelectExprs {
    82  		ae, ok := selectExpr.(*sqlparser.AliasedExpr)
    83  		if !ok {
    84  			return nil, vterrors.VT12001("pushing predicates on UNION where the first SELECT contains * or NEXT")
    85  		}
    86  		if !ae.As.IsEmpty() {
    87  			offsets[ae.As.String()] = i
    88  			continue
    89  		}
    90  		col, ok := ae.Expr.(*sqlparser.ColName)
    91  		if ok {
    92  			offsets[col.Name.Lowered()] = i
    93  		}
    94  	}
    95  
    96  	for i := range u.Sources {
    97  		var err error
    98  		predicate := sqlparser.CopyOnRewrite(expr, nil, func(cursor *sqlparser.CopyOnWriteCursor) {
    99  			col, ok := cursor.Node().(*sqlparser.ColName)
   100  			if !ok {
   101  				return
   102  			}
   103  
   104  			idx, ok := offsets[col.Name.Lowered()]
   105  			if !ok {
   106  				err = vterrors.VT13001("cannot push predicates on concatenate, missing columns from the UNION")
   107  				cursor.StopTreeWalk()
   108  				return
   109  			}
   110  
   111  			var sel *sqlparser.Select
   112  			sel, err = u.GetSelectFor(i)
   113  			if err != nil {
   114  				cursor.StopTreeWalk()
   115  				return
   116  			}
   117  
   118  			ae, ok := sel.SelectExprs[idx].(*sqlparser.AliasedExpr)
   119  			if !ok {
   120  				err = vterrors.VT12001("pushing non-aliased expression predicates on concatenate")
   121  				cursor.StopTreeWalk()
   122  				return
   123  			}
   124  			cursor.Replace(ae.Expr)
   125  		}, nil).(sqlparser.Expr)
   126  		if err != nil {
   127  			return nil, err
   128  		}
   129  		u.Sources[i], err = u.Sources[i].AddPredicate(ctx, predicate)
   130  		if err != nil {
   131  			return nil, err
   132  		}
   133  	}
   134  
   135  	return u, nil
   136  }
   137  
   138  func (u *Union) GetSelectFor(source int) (*sqlparser.Select, error) {
   139  	src := u.Sources[source]
   140  	for {
   141  		switch op := src.(type) {
   142  		case *Horizon:
   143  			return sqlparser.GetFirstSelect(op.Select), nil
   144  		case *Route:
   145  			src = op.Source
   146  		default:
   147  			return nil, vterrors.VT13001("expected all sources of the UNION to be horizons")
   148  		}
   149  	}
   150  }
   151  
   152  func (u *Union) Compact(*plancontext.PlanningContext) (ops.Operator, rewrite.TreeIdentity, error) {
   153  	var newSources []ops.Operator
   154  	anythingChanged := false
   155  	for _, source := range u.Sources {
   156  		var other *Union
   157  		horizon, ok := source.(*Horizon)
   158  		if ok {
   159  			union, ok := horizon.Source.(*Union)
   160  			if ok {
   161  				other = union
   162  			}
   163  		}
   164  		if other == nil {
   165  			newSources = append(newSources, source)
   166  			continue
   167  		}
   168  		anythingChanged = true
   169  		switch {
   170  		case len(other.Ordering) == 0 && !other.Distinct:
   171  			fallthrough
   172  		case u.Distinct:
   173  			// if the current UNION is a DISTINCT, we can safely ignore everything from children UNIONs, except LIMIT
   174  			newSources = append(newSources, other.Sources...)
   175  
   176  		default:
   177  			newSources = append(newSources, other)
   178  		}
   179  	}
   180  	if anythingChanged {
   181  		u.Sources = newSources
   182  	}
   183  	identity := rewrite.SameTree
   184  	if anythingChanged {
   185  		identity = rewrite.NewTree
   186  	}
   187  
   188  	return u, identity, nil
   189  }