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 }