From 08ac8fde23e42cf26677a3dfd1c3a0fb60e40f65 Mon Sep 17 00:00:00 2001 From: Marcus Gartner Date: Wed, 25 Jan 2023 18:47:32 -0500 Subject: [PATCH] sql: refactor tree.RoutineExpr `tree.RoutineExpr` no longer tracks the number of statements in the routine. Instead, it has a `tree.RoutinePlanGenerator` that generates a plan for each statement in the routine, given a list of arguments, and invokes a given callback on the plan. Release note: None --- pkg/sql/logictest/testdata/logic_test/udf | 8 +- pkg/sql/opt/exec/execbuilder/scalar.go | 204 ++++++++++++---------- pkg/sql/routine.go | 65 ++++--- pkg/sql/sem/tree/routine.go | 36 ++-- 4 files changed, 160 insertions(+), 153 deletions(-) diff --git a/pkg/sql/logictest/testdata/logic_test/udf b/pkg/sql/logictest/testdata/logic_test/udf index efcb3f6882f6..3fd4f4566818 100644 --- a/pkg/sql/logictest/testdata/logic_test/udf +++ b/pkg/sql/logictest/testdata/logic_test/udf @@ -2157,15 +2157,15 @@ SELECT trace_fn(a) FROM trace_tab statement ok SET tracing = off -query T +query T rowsort SELECT message FROM [SHOW TRACE FOR SESSION] WHERE message ~ 'udf' ---- -=== SPAN START: udf-stmt-trace_fn-0 === === SPAN START: udf-stmt-trace_fn-1 === -=== SPAN START: udf-stmt-trace_fn-0 === +=== SPAN START: udf-stmt-trace_fn-2 === === SPAN START: udf-stmt-trace_fn-1 === -=== SPAN START: udf-stmt-trace_fn-0 === +=== SPAN START: udf-stmt-trace_fn-2 === === SPAN START: udf-stmt-trace_fn-1 === +=== SPAN START: udf-stmt-trace_fn-2 === subtest args diff --git a/pkg/sql/opt/exec/execbuilder/scalar.go b/pkg/sql/opt/exec/execbuilder/scalar.go index 576697f428b7..0c6c0eab22bc 100644 --- a/pkg/sql/opt/exec/execbuilder/scalar.go +++ b/pkg/sql/opt/exec/execbuilder/scalar.go @@ -651,12 +651,11 @@ func (b *Builder) buildSubquery( // Create a tree.RoutinePlanFn that can plan the single statement // representing the subquery. - planFn := b.buildRoutinePlanFn(params, stmts, true /* allowOuterWithRefs */) + planGen := b.buildRoutinePlanGenerator(params, stmts, true /* allowOuterWithRefs */) return tree.NewTypedRoutineExpr( "subquery", args, - planFn, - 1, /* numStmts */ + planGen, subquery.Typ, false, /* enableStepping */ ), nil @@ -675,36 +674,43 @@ func (b *Builder) buildSubquery( inputRowCount := int64(input.Relational().Statistics().RowCountIfAvailable()) withExprs := make([]builtWithExpr, len(b.withExprs)) copy(withExprs, b.withExprs) - planFn := func( - ctx context.Context, ref tree.RoutineExecFactory, stmtIdx int, args tree.Datums, - ) (tree.RoutinePlan, error) { + planGen := func( + ctx context.Context, ref tree.RoutineExecFactory, args tree.Datums, fn tree.RoutinePlanGeneratedFunc, + ) error { ef := ref.(exec.Factory) eb := New(ctx, ef, b.optimizer, b.mem, b.catalog, input, b.evalCtx, false /* allowAutoCommit */, b.IsANSIDML) eb.withExprs = withExprs eb.disableTelemetry = true eb.planLazySubqueries = true - plan, err := eb.buildRelational(input) + ePlan, err := eb.buildRelational(input) if err != nil { - return nil, err + return err } if len(eb.subqueries) > 0 { - return nil, expectedLazyRoutineError("subquery") + return expectedLazyRoutineError("subquery") } if len(eb.cascades) > 0 { - return nil, expectedLazyRoutineError("cascade") + return expectedLazyRoutineError("cascade") } if len(eb.checks) > 0 { - return nil, expectedLazyRoutineError("check") + return expectedLazyRoutineError("check") } - return b.factory.ConstructPlan( - plan.root, nil /* subqueries */, nil /* cascades */, nil /* checks */, inputRowCount, + plan, err := b.factory.ConstructPlan( + ePlan.root, nil /* subqueries */, nil /* cascades */, nil /* checks */, inputRowCount, ) + if err != nil { + return err + } + err = fn(plan, true /* isFinalPlan */) + if err != nil { + return err + } + return nil } return tree.NewTypedRoutineExpr( "subquery", nil, /* args */ - planFn, - 1, /* numStmts */ + planGen, subquery.Typ, false, /* enableStepping */ ), nil @@ -770,7 +776,7 @@ func (b *Builder) buildUDF(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.Typ // Create a tree.RoutinePlanFn that can plan the statements in the UDF body. // TODO(mgartner): Add support for WITH expressions inside UDF bodies. - planFn := b.buildRoutinePlanFn(udf.Params, udf.Body, false /* allowOuterWithRefs */) + planGen := b.buildRoutinePlanGenerator(udf.Params, udf.Body, false /* allowOuterWithRefs */) // Enable stepping for volatile functions so that statements within the UDF // see mutations made by the invoking statement and by previous executed @@ -780,18 +786,17 @@ func (b *Builder) buildUDF(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.Typ return tree.NewTypedRoutineExpr( udf.Name, args, - planFn, - len(udf.Body), + planGen, udf.Typ, enableStepping, ), nil } -// buildRoutinePlanFn returns a tree.RoutinePlanFn that can plan the statements +// buildRoutinePlanGenerator returns a tree.RoutinePlanFn that can plan the statements // in a routine that has one or more arguments. -func (b *Builder) buildRoutinePlanFn( +func (b *Builder) buildRoutinePlanGenerator( params opt.ColList, stmts memo.RelListExpr, allowOuterWithRefs bool, -) tree.RoutinePlanFn { +) tree.RoutinePlanGenerator { // argOrd returns the ordinal of the argument within the arguments list that // can be substituted for each reference to the given function parameter // column. If the given column does not represent a function parameter, @@ -821,88 +826,95 @@ func (b *Builder) buildRoutinePlanFn( // // Note: we put o outside of the function so we allocate it only once. var o xform.Optimizer - planFn := func( - ctx context.Context, ref tree.RoutineExecFactory, stmtIdx int, args tree.Datums, - ) (tree.RoutinePlan, error) { - o.Init(ctx, b.evalCtx, b.catalog) - f := o.Factory() - stmt := stmts[stmtIdx] - - // Copy the expression into a new memo. Replace parameter references - // with argument datums. - addedWithBindings := false - var replaceFn norm.ReplaceFunc - replaceFn = func(e opt.Expr) opt.Expr { - switch t := e.(type) { - case *memo.VariableExpr: - if ord, ok := argOrd(t.Col); ok { - return f.ConstructConstVal(args[ord], t.Typ) - } - - case *memo.WithScanExpr: - // Allow referring to "outer" With expressions, if - // allowOuterWithRefs is true. The bound expressions are not - // part of this Memo, but they are used only for their - // relational properties, which should be valid. - // - // We must add all With expressions to the metadata even if they - // aren't referred to directly because they might be referred to - // transitively through other With expressions. For example, if - // stmt refers to With expression &1, and &1 refers to With - // expression &2, we must include &2 in the metadata so that its - // relational properties are available. See #87733. - // - // We lazily add these With expressions to the metadata here - // because the call to Factory.CopyAndReplace below clears With - // expressions in the metadata. - if allowOuterWithRefs && !addedWithBindings { - b.mem.Metadata().ForEachWithBinding(func(id opt.WithID, expr opt.Expr) { - f.Metadata().AddWithBinding(id, expr) - }) - addedWithBindings = true + planGen := func( + ctx context.Context, ref tree.RoutineExecFactory, args tree.Datums, fn tree.RoutinePlanGeneratedFunc, + ) error { + for i := range stmts { + stmt := stmts[i] + o.Init(ctx, b.evalCtx, b.catalog) + f := o.Factory() + + // Copy the expression into a new memo. Replace parameter references + // with argument datums. + addedWithBindings := false + var replaceFn norm.ReplaceFunc + replaceFn = func(e opt.Expr) opt.Expr { + switch t := e.(type) { + case *memo.VariableExpr: + if ord, ok := argOrd(t.Col); ok { + return f.ConstructConstVal(args[ord], t.Typ) + } + + case *memo.WithScanExpr: + // Allow referring to "outer" With expressions, if + // allowOuterWithRefs is true. The bound expressions are not + // part of this Memo, but they are used only for their + // relational properties, which should be valid. + // + // We must add all With expressions to the metadata even if they + // aren't referred to directly because they might be referred to + // transitively through other With expressions. For example, if + // stmt refers to With expression &1, and &1 refers to With + // expression &2, we must include &2 in the metadata so that its + // relational properties are available. See #87733. + // + // We lazily add these With expressions to the metadata here + // because the call to Factory.CopyAndReplace below clears With + // expressions in the metadata. + if allowOuterWithRefs && !addedWithBindings { + b.mem.Metadata().ForEachWithBinding(func(id opt.WithID, expr opt.Expr) { + f.Metadata().AddWithBinding(id, expr) + }) + addedWithBindings = true + } + // Fall through. } - // Fall through. + return f.CopyAndReplaceDefault(e, replaceFn) } - return f.CopyAndReplaceDefault(e, replaceFn) - } - f.CopyAndReplace(stmt, stmt.PhysProps, replaceFn) + f.CopyAndReplace(stmt, stmt.PhysProps, replaceFn) - // Optimize the memo. - optimizedExpr, err := o.Optimize() - if err != nil { - return nil, err - } + // Optimize the memo. + optimizedExpr, err := o.Optimize() + if err != nil { + return err + } - // Build the memo into a plan. - ef := ref.(exec.Factory) - eb := New(ctx, ef, &o, f.Memo(), b.catalog, optimizedExpr, b.evalCtx, false /* allowAutoCommit */, b.IsANSIDML) - eb.withExprs = withExprs - eb.disableTelemetry = true - eb.planLazySubqueries = true - plan, err := eb.Build() - if err != nil { - if errors.IsAssertionFailure(err) { - // Enhance the error with the EXPLAIN (OPT, VERBOSE) of the - // inner expression. - fmtFlags := memo.ExprFmtHideQualifications | memo.ExprFmtHideScalars | - memo.ExprFmtHideTypes - explainOpt := o.FormatExpr(optimizedExpr, fmtFlags) - err = errors.WithDetailf(err, "routineExpr:\n%s", explainOpt) + // Build the memo into a plan. + ef := ref.(exec.Factory) + eb := New(ctx, ef, &o, f.Memo(), b.catalog, optimizedExpr, b.evalCtx, false /* allowAutoCommit */, b.IsANSIDML) + eb.withExprs = withExprs + eb.disableTelemetry = true + eb.planLazySubqueries = true + plan, err := eb.Build() + if err != nil { + if errors.IsAssertionFailure(err) { + // Enhance the error with the EXPLAIN (OPT, VERBOSE) of the + // inner expression. + fmtFlags := memo.ExprFmtHideQualifications | memo.ExprFmtHideScalars | + memo.ExprFmtHideTypes + explainOpt := o.FormatExpr(optimizedExpr, fmtFlags) + err = errors.WithDetailf(err, "routineExpr:\n%s", explainOpt) + } + return err + } + if len(eb.subqueries) > 0 { + return expectedLazyRoutineError("subquery") + } + if len(eb.cascades) > 0 { + return expectedLazyRoutineError("cascade") + } + if len(eb.checks) > 0 { + return expectedLazyRoutineError("check") + } + isFinalPlan := i == len(stmts)-1 + err = fn(plan, isFinalPlan) + if err != nil { + return err } - return nil, err - } - if len(eb.subqueries) > 0 { - return nil, expectedLazyRoutineError("subquery") - } - if len(eb.cascades) > 0 { - return nil, expectedLazyRoutineError("cascade") - } - if len(eb.checks) > 0 { - return nil, expectedLazyRoutineError("check") } - return plan, nil + return nil } - return planFn + return planGen } func expectedLazyRoutineError(typ string) error { diff --git a/pkg/sql/routine.go b/pkg/sql/routine.go index 8cc833fdfc65..8f02858a03fa 100644 --- a/pkg/sql/routine.go +++ b/pkg/sql/routine.go @@ -25,7 +25,7 @@ import ( // runs the plans. The resulting value of the last statement in the routine is // returned. func (p *planner) EvalRoutineExpr( - ctx context.Context, expr *tree.RoutineExpr, input tree.Datums, + ctx context.Context, expr *tree.RoutineExpr, args tree.Datums, ) (result tree.Datum, err error) { // Return the cached result if it exists. if expr.CachedResult != nil { @@ -59,46 +59,41 @@ func (p *planner) EvalRoutineExpr( } // Execute each statement in the routine sequentially. + stmtIdx := 0 ef := newExecFactory(ctx, p) - for i := 0; i < expr.NumStmts; i++ { - if err := func() error { - opName := "udf-stmt-" + expr.Name + "-" + strconv.Itoa(i) - ctx, sp := tracing.ChildSpan(ctx, opName) - defer sp.Finish() - - // Generate a plan for executing the ith statement. - plan, err := expr.PlanFn(ctx, ef, i, input) - if err != nil { - return err - } - - // If this is the last statement, use the rowResultWriter created above. - // Otherwise, use a rowResultWriter that drops all rows added to it. - var w rowResultWriter - if i == expr.NumStmts-1 { - w = rrw - } else { - w = &droppingResultWriter{} - } - - // Place a sequence point before each statement in the routine for - // volatile functions. - if expr.EnableStepping { - if err := txn.Step(ctx); err != nil { - return err - } - } + err = expr.ForEachPlan(ctx, ef, args, func(plan tree.RoutinePlan, isFinalPlan bool) error { + stmtIdx++ + opName := "udf-stmt-" + expr.Name + "-" + strconv.Itoa(stmtIdx) + ctx, sp := tracing.ChildSpan(ctx, opName) + defer sp.Finish() + + // If this is the last statement, use the rowResultWriter created above. + // Otherwise, use a rowResultWriter that drops all rows added to it. + var w rowResultWriter + if isFinalPlan { + w = rrw + } else { + w = &droppingResultWriter{} + } - // Run the plan. - err = runPlanInsidePlan(ctx, p.RunParams(ctx), plan.(*planComponents), w) - if err != nil { + // Place a sequence point before each statement in the routine for + // volatile functions. + if expr.EnableStepping { + if err := txn.Step(ctx); err != nil { return err } + } - return nil - }(); err != nil { - return nil, err + // Run the plan. + err = runPlanInsidePlan(ctx, p.RunParams(ctx), plan.(*planComponents), w) + if err != nil { + return err } + + return nil + }) + if err != nil { + return nil, err } // Fetch the first row from the row container and return the first diff --git a/pkg/sql/sem/tree/routine.go b/pkg/sql/sem/tree/routine.go index 904518a6bfd4..3dae28b4d8dc 100644 --- a/pkg/sql/sem/tree/routine.go +++ b/pkg/sql/sem/tree/routine.go @@ -16,11 +16,20 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/types" ) -// RoutinePlanFn creates a plan for the execution of one statement within a -// routine. -type RoutinePlanFn func( - _ context.Context, _ RoutineExecFactory, stmtIdx int, args Datums, -) (RoutinePlan, error) +// RoutinePlanGenerator generates a plan for the execution of each statement +// within a routine. The given RoutinePlanGeneratedFunc is called for each plan +// generated. +// +// A RoutinePlanGenerator must return an error if the RoutinePlanGeneratedFunc +// returns an error. +type RoutinePlanGenerator func( + _ context.Context, _ RoutineExecFactory, args Datums, fn RoutinePlanGeneratedFunc, +) error + +// RoutinePlanGeneratedFunc is the function type that is called for each plan +// enumerated by a RoutinePlanGenerator. isFinalPlan is true if no more plans +// will be generated after the current plan. +type RoutinePlanGeneratedFunc func(plan RoutinePlan, isFinalPlan bool) error // RoutinePlan represents a plan for a statement in a routine. It currently maps // to exec.Plan. We use the empty interface here rather then exec.Plan to avoid @@ -45,11 +54,8 @@ type RoutineExpr struct { // Args contains the argument expressions to the routine. Args TypedExprs - // PlanFn returns an exec plan for a given statement in the routine. - PlanFn RoutinePlanFn - - // NumStmts is the number of statements in the routine. - NumStmts int + // ForEachPlan generates a plan for each statement in the routine. + ForEachPlan RoutinePlanGenerator // Typ is the type of the routine's result. Typ *types.T @@ -91,17 +97,11 @@ type RoutineExpr struct { // NewTypedRoutineExpr returns a new RoutineExpr that is well-typed. func NewTypedRoutineExpr( - name string, - args TypedExprs, - planFn RoutinePlanFn, - numStmts int, - typ *types.T, - enableStepping bool, + name string, args TypedExprs, gen RoutinePlanGenerator, typ *types.T, enableStepping bool, ) *RoutineExpr { return &RoutineExpr{ Args: args, - PlanFn: planFn, - NumStmts: numStmts, + ForEachPlan: gen, Typ: typ, EnableStepping: enableStepping, Name: name,