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

     1  // Copyright 2016 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 rowexec
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    21  	"github.com/cockroachdb/cockroach/pkg/util/mon"
    22  	"github.com/cockroachdb/errors"
    23  )
    24  
    25  // streamGroupAccumulator groups input rows coming from src into groups dictated
    26  // by equality according to the ordering columns.
    27  type streamGroupAccumulator struct {
    28  	src   execinfra.RowSource
    29  	types []*types.T
    30  
    31  	// srcConsumed is set once src has been exhausted.
    32  	srcConsumed bool
    33  	ordering    sqlbase.ColumnOrdering
    34  
    35  	// curGroup maintains the rows accumulated in the current group.
    36  	curGroup   []sqlbase.EncDatumRow
    37  	datumAlloc sqlbase.DatumAlloc
    38  
    39  	// leftoverRow is the first row of the next group. It's saved in the
    40  	// accumulator after the current group is returned, so the accumulator can
    41  	// resume later.
    42  	leftoverRow sqlbase.EncDatumRow
    43  
    44  	rowAlloc sqlbase.EncDatumRowAlloc
    45  
    46  	memAcc mon.BoundAccount
    47  }
    48  
    49  func makeStreamGroupAccumulator(
    50  	src execinfra.RowSource, ordering sqlbase.ColumnOrdering, memMonitor *mon.BytesMonitor,
    51  ) streamGroupAccumulator {
    52  	return streamGroupAccumulator{
    53  		src:      src,
    54  		types:    src.OutputTypes(),
    55  		ordering: ordering,
    56  		memAcc:   memMonitor.MakeBoundAccount(),
    57  	}
    58  }
    59  
    60  func (s *streamGroupAccumulator) start(ctx context.Context) {
    61  	s.src.Start(ctx)
    62  }
    63  
    64  // nextGroup returns the next group from the inputs. The returned slice is not safe
    65  // to use after the next call to nextGroup.
    66  func (s *streamGroupAccumulator) nextGroup(
    67  	ctx context.Context, evalCtx *tree.EvalContext,
    68  ) ([]sqlbase.EncDatumRow, *execinfrapb.ProducerMetadata) {
    69  	if s.srcConsumed {
    70  		// If src has been exhausted, then we also must have advanced away from the
    71  		// last group.
    72  		return nil, nil
    73  	}
    74  
    75  	if s.leftoverRow != nil {
    76  		s.curGroup = append(s.curGroup, s.leftoverRow)
    77  		s.leftoverRow = nil
    78  	}
    79  
    80  	for {
    81  		row, meta := s.src.Next()
    82  		if meta != nil {
    83  			return nil, meta
    84  		}
    85  		if row == nil {
    86  			s.srcConsumed = true
    87  			return s.curGroup, nil
    88  		}
    89  
    90  		if err := s.memAcc.Grow(ctx, int64(row.Size())); err != nil {
    91  			return nil, &execinfrapb.ProducerMetadata{Err: err}
    92  		}
    93  		row = s.rowAlloc.CopyRow(row)
    94  
    95  		if len(s.curGroup) == 0 {
    96  			if s.curGroup == nil {
    97  				s.curGroup = make([]sqlbase.EncDatumRow, 0, 64)
    98  			}
    99  			s.curGroup = append(s.curGroup, row)
   100  			continue
   101  		}
   102  
   103  		cmp, err := s.curGroup[0].Compare(s.types, &s.datumAlloc, s.ordering, evalCtx, row)
   104  		if err != nil {
   105  			return nil, &execinfrapb.ProducerMetadata{Err: err}
   106  		}
   107  		if cmp == 0 {
   108  			s.curGroup = append(s.curGroup, row)
   109  		} else if cmp == 1 {
   110  			return nil, &execinfrapb.ProducerMetadata{
   111  				Err: errors.Errorf(
   112  					"detected badly ordered input: %s > %s, but expected '<'",
   113  					s.curGroup[0].String(s.types), row.String(s.types)),
   114  			}
   115  		} else {
   116  			n := len(s.curGroup)
   117  			ret := s.curGroup[:n:n]
   118  			s.curGroup = s.curGroup[:0]
   119  			s.memAcc.Empty(ctx)
   120  			s.leftoverRow = row
   121  			return ret, nil
   122  		}
   123  	}
   124  }
   125  
   126  func (s *streamGroupAccumulator) close(ctx context.Context) {
   127  	s.memAcc.Close(ctx)
   128  }