github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/opt/norm/with_funcs.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package norm
    12  
    13  import (
    14  	"github.com/cockroachdb/cockroach/pkg/sql/opt"
    15  	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
    16  	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
    17  )
    18  
    19  // CanInlineWith returns whether or not it's valid to inline binding in expr.
    20  // This is the case when materialize is explicitly set to false, or when:
    21  // 1. binding has no side-effects (because once it's inlined, there's no
    22  //    guarantee it will be executed fully), and
    23  // 2. binding is referenced at most once in expr.
    24  func (c *CustomFuncs) CanInlineWith(binding, expr memo.RelExpr, private *memo.WithPrivate) bool {
    25  	// If materialization is set, ignore the checks below.
    26  	if private.Mtr.Set {
    27  		return !private.Mtr.Materialize
    28  	}
    29  	if binding.Relational().CanHaveSideEffects {
    30  		return false
    31  	}
    32  	return c.WithUses(expr)[private.ID].Count <= 1
    33  }
    34  
    35  // InlineWith replaces all references to the With expression in input (via
    36  // WithScans) with its definition.
    37  func (c *CustomFuncs) InlineWith(binding, input memo.RelExpr, priv *memo.WithPrivate) memo.RelExpr {
    38  	var replace ReplaceFunc
    39  	replace = func(nd opt.Expr) opt.Expr {
    40  		switch t := nd.(type) {
    41  		case *memo.WithScanExpr:
    42  			if t.With == priv.ID {
    43  				// TODO(justin): it might be worth carefully walking the tree and
    44  				// renaming variables as we do this replacement so that this projection
    45  				// is unnecessary (assuming there's at most one reference to the
    46  				// WithScan, which might be false if we heuristically inline multiple
    47  				// times in the future).
    48  				projections := make(memo.ProjectionsExpr, len(t.InCols))
    49  				for i := range t.InCols {
    50  					projections[i] = c.f.ConstructProjectionsItem(
    51  						c.f.ConstructVariable(t.InCols[i]),
    52  						t.OutCols[i],
    53  					)
    54  				}
    55  				return c.f.ConstructProject(binding, projections, opt.ColSet{})
    56  			}
    57  			// TODO(justin): should apply joins block inlining because they can lead
    58  			// to expressions being executed multiple times?
    59  		}
    60  		return c.f.Replace(nd, replace)
    61  	}
    62  
    63  	return replace(input).(memo.RelExpr)
    64  }
    65  
    66  // WithUses returns the WithUsesMap for the given expression.
    67  func (c *CustomFuncs) WithUses(r opt.Expr) props.WithUsesMap {
    68  	switch e := r.(type) {
    69  	case memo.RelExpr:
    70  		relProps := e.Relational()
    71  
    72  		// Lazily calculate and store the WithUses value.
    73  		if !relProps.IsAvailable(props.WithUses) {
    74  			relProps.Shared.Rule.WithUses = c.deriveWithUses(r)
    75  			relProps.SetAvailable(props.WithUses)
    76  		}
    77  		return relProps.Shared.Rule.WithUses
    78  
    79  	case memo.ScalarPropsExpr:
    80  		scalarProps := e.ScalarProps()
    81  
    82  		// Lazily calculate and store the WithUses value.
    83  		if !scalarProps.IsAvailable(props.WithUses) {
    84  			scalarProps.Shared.Rule.WithUses = c.deriveWithUses(r)
    85  			scalarProps.SetAvailable(props.WithUses)
    86  		}
    87  		return scalarProps.Shared.Rule.WithUses
    88  
    89  	default:
    90  		return c.deriveWithUses(r)
    91  	}
    92  }
    93  
    94  // deriveWithUses collects information about WithScans in the expression.
    95  func (c *CustomFuncs) deriveWithUses(r opt.Expr) props.WithUsesMap {
    96  	// We don't allow the information to escape the scope of the WITH itself, so
    97  	// we exclude that ID from the results.
    98  	var excludedID opt.WithID
    99  
   100  	switch e := r.(type) {
   101  	case *memo.WithScanExpr:
   102  		info := props.WithUseInfo{
   103  			Count:    1,
   104  			UsedCols: e.InCols.ToSet(),
   105  		}
   106  		return props.WithUsesMap{e.With: info}
   107  
   108  	case *memo.WithExpr:
   109  		excludedID = e.ID
   110  
   111  	default:
   112  		if opt.IsMutationOp(e) {
   113  			// Note: this can still be 0.
   114  			excludedID = e.Private().(*memo.MutationPrivate).WithID
   115  		}
   116  	}
   117  
   118  	var result props.WithUsesMap
   119  	for i, n := 0, r.ChildCount(); i < n; i++ {
   120  		childUses := c.WithUses(r.Child(i))
   121  		for id, info := range childUses {
   122  			if id == excludedID {
   123  				continue
   124  			}
   125  			if result == nil {
   126  				result = make(props.WithUsesMap, len(childUses))
   127  			}
   128  			existing := result[id]
   129  			existing.Count += info.Count
   130  			existing.UsedCols.UnionWith(info.UsedCols)
   131  			result[id] = existing
   132  		}
   133  	}
   134  	return result
   135  }