github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/rowexec/countrows.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 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/tracing"
    22  	"github.com/opentracing/opentracing-go"
    23  )
    24  
    25  // countAggregator is a simple processor that counts the number of rows it
    26  // receives. It's a specialized aggregator that can be used for COUNT(*).
    27  type countAggregator struct {
    28  	execinfra.ProcessorBase
    29  
    30  	input execinfra.RowSource
    31  	count int
    32  }
    33  
    34  var _ execinfra.Processor = &countAggregator{}
    35  var _ execinfra.RowSource = &countAggregator{}
    36  
    37  const countRowsProcName = "count rows"
    38  
    39  var outputTypes = []*types.T{types.Int}
    40  
    41  func newCountAggregator(
    42  	flowCtx *execinfra.FlowCtx,
    43  	processorID int32,
    44  	input execinfra.RowSource,
    45  	post *execinfrapb.PostProcessSpec,
    46  	output execinfra.RowReceiver,
    47  ) (*countAggregator, error) {
    48  	ag := &countAggregator{}
    49  	ag.input = input
    50  
    51  	if sp := opentracing.SpanFromContext(flowCtx.EvalCtx.Ctx()); sp != nil && tracing.IsRecording(sp) {
    52  		ag.input = newInputStatCollector(input)
    53  		ag.FinishTrace = ag.outputStatsToTrace
    54  	}
    55  
    56  	if err := ag.Init(
    57  		ag,
    58  		post,
    59  		outputTypes,
    60  		flowCtx,
    61  		processorID,
    62  		output,
    63  		nil, /* memMonitor */
    64  		execinfra.ProcStateOpts{
    65  			InputsToDrain: []execinfra.RowSource{ag.input},
    66  		},
    67  	); err != nil {
    68  		return nil, err
    69  	}
    70  
    71  	return ag, nil
    72  }
    73  
    74  func (ag *countAggregator) Start(ctx context.Context) context.Context {
    75  	ag.input.Start(ctx)
    76  	return ag.StartInternal(ctx, countRowsProcName)
    77  }
    78  
    79  func (ag *countAggregator) Next() (sqlbase.EncDatumRow, *execinfrapb.ProducerMetadata) {
    80  	for ag.State == execinfra.StateRunning {
    81  		row, meta := ag.input.Next()
    82  		if meta != nil {
    83  			if meta.Err != nil {
    84  				ag.MoveToDraining(meta.Err)
    85  				break
    86  			}
    87  			return nil, meta
    88  		}
    89  		if row == nil {
    90  			ret := make(sqlbase.EncDatumRow, 1)
    91  			ret[0] = sqlbase.EncDatum{Datum: tree.NewDInt(tree.DInt(ag.count))}
    92  			rendered, _, err := ag.Out.ProcessRow(ag.Ctx, ret)
    93  			// We're done as soon as we process our one output row, so we
    94  			// transition into draining state. We will, however, return non-nil
    95  			// error (if such occurs during rendering) separately below.
    96  			ag.MoveToDraining(nil /* err */)
    97  			if err != nil {
    98  				return nil, &execinfrapb.ProducerMetadata{Err: err}
    99  			}
   100  			return rendered, nil
   101  		}
   102  		ag.count++
   103  	}
   104  	return nil, ag.DrainHelper()
   105  }
   106  
   107  func (ag *countAggregator) ConsumerDone() {
   108  	ag.MoveToDraining(nil /* err */)
   109  }
   110  
   111  func (ag *countAggregator) ConsumerClosed() {
   112  	ag.InternalClose()
   113  }
   114  
   115  // outputStatsToTrace outputs the collected distinct stats to the trace. Will
   116  // fail silently if the Distinct processor is not collecting stats.
   117  func (ag *countAggregator) outputStatsToTrace() {
   118  	is, ok := getInputStats(ag.FlowCtx, ag.input)
   119  	if !ok {
   120  		return
   121  	}
   122  	if sp := opentracing.SpanFromContext(ag.Ctx); sp != nil {
   123  		tracing.SetSpanStats(
   124  			sp, &AggregatorStats{InputStats: is},
   125  		)
   126  	}
   127  }