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

     1  // Copyright 2017 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  // This file contains helper code to populate execinfrapb.Expressions during
    12  // planning.
    13  
    14  package physicalplan
    15  
    16  import (
    17  	"fmt"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    23  	"github.com/cockroachdb/cockroach/pkg/util/log"
    24  )
    25  
    26  // ExprContext is an interface containing objects necessary for creating
    27  // execinfrapb.Expressions.
    28  type ExprContext interface {
    29  	// EvalContext returns the tree.EvalContext for planning.
    30  	EvalContext() *tree.EvalContext
    31  
    32  	// IsLocal returns true if the current plan is local.
    33  	IsLocal() bool
    34  
    35  	// EvaluateSubqueries returns true if subqueries should be evaluated before
    36  	// creating the execinfrapb.Expression.
    37  	EvaluateSubqueries() bool
    38  }
    39  
    40  // fakeExprContext is a fake implementation of ExprContext that always behaves
    41  // as if it were part of a non-local query.
    42  type fakeExprContext struct{}
    43  
    44  var _ ExprContext = fakeExprContext{}
    45  
    46  func (fakeExprContext) EvalContext() *tree.EvalContext {
    47  	return &tree.EvalContext{}
    48  }
    49  
    50  func (fakeExprContext) IsLocal() bool {
    51  	return false
    52  }
    53  
    54  func (fakeExprContext) EvaluateSubqueries() bool {
    55  	return true
    56  }
    57  
    58  // MakeExpression creates a execinfrapb.Expression.
    59  //
    60  // The execinfrapb.Expression uses the placeholder syntax (@1, @2, @3..) to
    61  // refer to columns.
    62  //
    63  // The expr uses IndexedVars to refer to columns. The caller can optionally
    64  // remap these columns by passing an indexVarMap: an IndexedVar with index i
    65  // becomes column indexVarMap[i].
    66  //
    67  // ctx can be nil in which case a fakeExprCtx will be used.
    68  func MakeExpression(
    69  	expr tree.TypedExpr, ctx ExprContext, indexVarMap []int,
    70  ) (execinfrapb.Expression, error) {
    71  	if expr == nil {
    72  		return execinfrapb.Expression{}, nil
    73  	}
    74  	if ctx == nil {
    75  		ctx = &fakeExprContext{}
    76  	}
    77  
    78  	if ctx.IsLocal() {
    79  		if indexVarMap != nil {
    80  			// Remap our indexed vars.
    81  			expr = sqlbase.RemapIVarsInTypedExpr(expr, indexVarMap)
    82  		}
    83  		return execinfrapb.Expression{LocalExpr: expr}, nil
    84  	}
    85  
    86  	evalCtx := ctx.EvalContext()
    87  	subqueryVisitor := &evalAndReplaceSubqueryVisitor{
    88  		evalCtx: evalCtx,
    89  	}
    90  
    91  	outExpr := expr.(tree.Expr)
    92  	if ctx.EvaluateSubqueries() {
    93  		outExpr, _ = tree.WalkExpr(subqueryVisitor, expr)
    94  		if subqueryVisitor.err != nil {
    95  			return execinfrapb.Expression{}, subqueryVisitor.err
    96  		}
    97  	}
    98  	// We format the expression using the IndexedVar and Placeholder formatting interceptors.
    99  	fmtCtx := execinfrapb.ExprFmtCtxBase(evalCtx)
   100  	if indexVarMap != nil {
   101  		fmtCtx.SetIndexedVarFormat(func(ctx *tree.FmtCtx, idx int) {
   102  			remappedIdx := indexVarMap[idx]
   103  			if remappedIdx < 0 {
   104  				panic(fmt.Sprintf("unmapped index %d", idx))
   105  			}
   106  			ctx.Printf("@%d", remappedIdx+1)
   107  		})
   108  	}
   109  	fmtCtx.FormatNode(outExpr)
   110  	if log.V(1) {
   111  		log.Infof(evalCtx.Ctx(), "Expr %s:\n%s", fmtCtx.String(), tree.ExprDebugString(outExpr))
   112  	}
   113  	return execinfrapb.Expression{Expr: fmtCtx.CloseAndGetString()}, nil
   114  }
   115  
   116  type evalAndReplaceSubqueryVisitor struct {
   117  	evalCtx *tree.EvalContext
   118  	err     error
   119  }
   120  
   121  var _ tree.Visitor = &evalAndReplaceSubqueryVisitor{}
   122  
   123  func (e *evalAndReplaceSubqueryVisitor) VisitPre(expr tree.Expr) (bool, tree.Expr) {
   124  	switch expr := expr.(type) {
   125  	case *tree.Subquery:
   126  		val, err := e.evalCtx.Planner.EvalSubquery(expr)
   127  		if err != nil {
   128  			e.err = err
   129  			return false, expr
   130  		}
   131  		var newExpr tree.Expr = val
   132  		if _, isTuple := val.(*tree.DTuple); !isTuple && expr.ResolvedType().Family() != types.UnknownFamily {
   133  			newExpr = tree.NewTypedCastExpr(val, expr.ResolvedType())
   134  		}
   135  		return false, newExpr
   136  	default:
   137  		return true, expr
   138  	}
   139  }
   140  
   141  func (evalAndReplaceSubqueryVisitor) VisitPost(expr tree.Expr) tree.Expr { return expr }