github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/recursive_cte.go (about) 1 // Copyright 2019 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 sql 12 13 import ( 14 "context" 15 16 "github.com/cockroachdb/cockroach/pkg/sql/opt/exec" 17 "github.com/cockroachdb/cockroach/pkg/sql/rowcontainer" 18 "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" 19 "github.com/cockroachdb/cockroach/pkg/sql/sqlbase" 20 ) 21 22 // recursiveCTENode implements the logic for a recursive CTE: 23 // 1. Evaluate the initial query; emit the results and also save them in 24 // a "working" table. 25 // 2. So long as the working table is not empty: 26 // * evaluate the recursive query, substituting the current contents of 27 // the working table for the recursive self-reference; 28 // * emit all resulting rows, and save them as the next iteration's 29 // working table. 30 // The recursive query tree is regenerated each time using a callback 31 // (implemented by the execbuilder). 32 type recursiveCTENode struct { 33 initial planNode 34 35 genIterationFn exec.RecursiveCTEIterationFn 36 37 label string 38 39 recursiveCTERun 40 } 41 42 type recursiveCTERun struct { 43 // workingRows contains the rows produced by the current iteration (aka the 44 // "working" table). 45 workingRows *rowcontainer.RowContainer 46 // nextRowIdx is the index inside workingRows of the next row to be returned 47 // by the operator. 48 nextRowIdx int 49 50 initialDone bool 51 done bool 52 } 53 54 func (n *recursiveCTENode) startExec(params runParams) error { 55 n.workingRows = rowcontainer.NewRowContainer( 56 params.EvalContext().Mon.MakeBoundAccount(), 57 sqlbase.ColTypeInfoFromResCols(getPlanColumns(n.initial, false /* mut */)), 58 0, /* rowCapacity */ 59 ) 60 n.nextRowIdx = 0 61 return nil 62 } 63 64 func (n *recursiveCTENode) Next(params runParams) (bool, error) { 65 if err := params.p.cancelChecker.Check(); err != nil { 66 return false, err 67 } 68 69 n.nextRowIdx++ 70 71 if !n.initialDone { 72 ok, err := n.initial.Next(params) 73 if err != nil { 74 return false, err 75 } 76 if ok { 77 if _, err = n.workingRows.AddRow(params.ctx, n.initial.Values()); err != nil { 78 return false, err 79 } 80 return true, nil 81 } 82 n.initialDone = true 83 } 84 85 if n.done { 86 return false, nil 87 } 88 89 if n.workingRows.Len() == 0 { 90 // Last iteration returned no rows. 91 n.done = true 92 return false, nil 93 } 94 95 // There are more rows to return from the last iteration. 96 if n.nextRowIdx <= n.workingRows.Len() { 97 return true, nil 98 } 99 100 // Let's run another iteration. 101 102 lastWorkingRows := n.workingRows 103 defer lastWorkingRows.Close(params.ctx) 104 105 n.workingRows = rowcontainer.NewRowContainer( 106 params.EvalContext().Mon.MakeBoundAccount(), 107 sqlbase.ColTypeInfoFromResCols(getPlanColumns(n.initial, false /* mut */)), 108 0, /* rowCapacity */ 109 ) 110 111 // Set up a bufferNode that can be used as a reference for a scanBufferNode. 112 buf := &bufferNode{ 113 // The plan here is only useful for planColumns, so it's ok to always use 114 // the initial plan. 115 plan: n.initial, 116 bufferedRows: lastWorkingRows, 117 label: n.label, 118 } 119 newPlan, err := n.genIterationFn(buf) 120 if err != nil { 121 return false, err 122 } 123 124 if err := runPlanInsidePlan(params, newPlan.(*planTop), n.workingRows); err != nil { 125 return false, err 126 } 127 n.nextRowIdx = 1 128 return n.workingRows.Len() > 0, nil 129 } 130 131 func (n *recursiveCTENode) Values() tree.Datums { 132 return n.workingRows.At(n.nextRowIdx - 1) 133 } 134 135 func (n *recursiveCTENode) Close(ctx context.Context) { 136 n.initial.Close(ctx) 137 n.workingRows.Close(ctx) 138 }