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

     1  // Copyright 2018 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/rowcontainer"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    19  )
    20  
    21  // spoolNode ensures that a child planNode is executed to completion
    22  // during the start phase. The results, if any, are collected. The
    23  // child node is guaranteed to run to completion.
    24  // If hardLimit is set, only that number of rows is collected, but
    25  // the child node is still run to completion.
    26  type spoolNode struct {
    27  	source    planNode
    28  	rows      *rowcontainer.RowContainer
    29  	hardLimit int64
    30  	curRowIdx int
    31  }
    32  
    33  func (s *spoolNode) startExec(params runParams) error {
    34  	// If FastPathResults() on the source indicates that the results are
    35  	// already available (2nd value true), then the computation is
    36  	// already done at start time and spooling is unnecessary.
    37  	if f, ok := s.source.(planNodeFastPath); ok {
    38  		_, done := f.FastPathResults()
    39  		if done {
    40  			return nil
    41  		}
    42  	}
    43  
    44  	s.rows = rowcontainer.NewRowContainer(
    45  		params.EvalContext().Mon.MakeBoundAccount(),
    46  		sqlbase.ColTypeInfoFromResCols(planColumns(s.source)),
    47  		0, /* rowCapacity */
    48  	)
    49  
    50  	// Accumulate all the rows up to the hardLimit, if any.
    51  	// This also guarantees execution of the child node to completion,
    52  	// even if Next() on the spool itself is not called for every row.
    53  	for {
    54  		next, err := s.source.Next(params)
    55  		if err != nil {
    56  			return err
    57  		}
    58  		if !next {
    59  			break
    60  		}
    61  		if s.hardLimit == 0 || int64(s.rows.Len()) < s.hardLimit {
    62  			if _, err := s.rows.AddRow(params.ctx, s.source.Values()); err != nil {
    63  				return err
    64  			}
    65  		}
    66  	}
    67  	s.curRowIdx = -1
    68  	return nil
    69  }
    70  
    71  // FastPathResults implements the planNodeFastPath interface.
    72  func (s *spoolNode) FastPathResults() (int, bool) {
    73  	// If the source implements the fast path interface, let it report
    74  	// its status through. This lets e.g. the fast path of a DELETE or
    75  	// an UPSERT report that they have finished its computing already,
    76  	// so the calls to Next() on the spool itself can also be elided.
    77  	// If FastPathResults() on the source says the fast path is unavailable,
    78  	// then startExec() on the spool will also notice that and
    79  	// spooling will occur as expected.
    80  	if f, ok := s.source.(planNodeFastPath); ok {
    81  		return f.FastPathResults()
    82  	}
    83  	return 0, false
    84  }
    85  
    86  // spooled implements the planNodeSpooled interface.
    87  func (s *spoolNode) spooled() {}
    88  
    89  // Next is part of the planNode interface.
    90  func (s *spoolNode) Next(params runParams) (bool, error) {
    91  	s.curRowIdx++
    92  	return s.curRowIdx < s.rows.Len(), nil
    93  }
    94  
    95  // Values is part of the planNode interface.
    96  func (s *spoolNode) Values() tree.Datums {
    97  	return s.rows.At(s.curRowIdx)
    98  }
    99  
   100  // Close is part of the planNode interface.
   101  func (s *spoolNode) Close(ctx context.Context) {
   102  	s.source.Close(ctx)
   103  	if s.rows != nil {
   104  		s.rows.Close(ctx)
   105  		s.rows = nil
   106  	}
   107  }