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  }