colexec: clean up meta components handling

pkg/sql/colexec/BUILD.bazel

@@ -31,7 +31,6 @@ go_library(
         "sort_chunks.go",
         "sort_utils.go",
         "sorttopk.go",
-        "stats.go",
         "tuple_proj_op.go",
         "unordered_distinct.go",
         "utils.go",

pkg/sql/colexec/colbuilder/BUILD.bazel

@@ -57,7 +57,6 @@ go_test(
         "//pkg/sql/catalog/catalogkv",
         "//pkg/sql/colexec",
         "//pkg/sql/colexec/colexecargs",
-        "//pkg/sql/colexecop",
         "//pkg/sql/execinfra",
         "//pkg/sql/execinfrapb",
         "//pkg/sql/randgen",

pkg/sql/colexec/colbuilder/execplan_test.go

@@ -21,7 +21,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkv"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecargs"
-	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/randgen"
@@ -114,7 +113,7 @@ func TestNewColOperatorExpectedTypeSchema(t *testing.T) {
 			Post:        execinfrapb.PostProcessSpec{RenderExprs: []execinfrapb.Expression{{Expr: "@1 - 1"}}},
 			ResultTypes: []*types.T{types.Int},
 		},
-		Inputs:              []colexecop.Operator{r.Op},
+		Inputs:              []colexecargs.OpWithMetaInfo{{Root: r.Root}},
 		StreamingMemAccount: &streamingMemAcc,
 	}
 	r, err = NewColOperator(ctx, flowCtx, args)
@@ -123,12 +122,9 @@ func TestNewColOperatorExpectedTypeSchema(t *testing.T) {
 	m, err := colexec.NewMaterializer(
 		flowCtx,
 		0, /* processorID */
-		r.Op,
+		r.OpWithMetaInfo,
 		[]*types.T{types.Int},
 		nil, /* output */
-		nil, /* getStats */
-		nil, /* metadataSources */
-		nil, /* toClose */
 		nil, /* cancelFlow */
 	)
 	require.NoError(t, err)

pkg/sql/colexec/colexecargs/op_creation.go

@@ -32,25 +32,37 @@ import (
 var TestNewColOperator func(ctx context.Context, flowCtx *execinfra.FlowCtx, args *NewColOperatorArgs,
 ) (r *NewColOperatorResult, err error)
 
+// OpWithMetaInfo stores a colexecop.Operator together with miscellaneous meta
+// information about the tree rooted in that operator.
+// TODO(yuzefovich): figure out the story about pooling these objects.
+type OpWithMetaInfo struct {
+	Root colexecop.Operator
+	// StatsCollectors are all stats collectors that are present in the tree
+	// rooted in Root for which the responsibility of retrieving stats hasn't
+	// been claimed yet.
+	StatsCollectors []colexecop.VectorizedStatsCollector
+	// MetadataSources are all sources of the metadata that are present in the
+	// tree rooted in Root for which the responsibility of draining hasn't been
+	// claimed yet.
+	MetadataSources colexecop.MetadataSources
+	// ToClose are all colexecop.Closers that are present in the tree rooted in
+	// Root for which the responsibility of closing hasn't been claimed yet.
+	ToClose colexecop.Closers
+}
+
 // NewColOperatorArgs is a helper struct that encompasses all of the input
 // arguments to NewColOperator call.
 type NewColOperatorArgs struct {
 	Spec                 *execinfrapb.ProcessorSpec
-	Inputs               []colexecop.Operator
+	Inputs               []OpWithMetaInfo
 	StreamingMemAccount  *mon.BoundAccount
 	ProcessorConstructor execinfra.ProcessorConstructor
 	LocalProcessors      []execinfra.LocalProcessor
-	// MetadataSources are all sources of the metadata that are present in the
-	// trees rooted in Inputs. The slice has the same length as Inputs. If
-	// NewColOperator call creates a colexec.Materializer, then it will take
-	// over the responsibility of draining the sources; otherwise, they will be
-	// returned in NewColOperatorResult.
-	MetadataSources []colexecop.MetadataSources
-	DiskQueueCfg    colcontainer.DiskQueueCfg
-	FDSemaphore     semaphore.Semaphore
-	ExprHelper      *ExprHelper
-	Factory         coldata.ColumnFactory
-	TestingKnobs    struct {
+	DiskQueueCfg         colcontainer.DiskQueueCfg
+	FDSemaphore          semaphore.Semaphore
+	ExprHelper           *ExprHelper
+	Factory              coldata.ColumnFactory
+	TestingKnobs         struct {
 		// SpillingCallbackFn will be called when the spilling from an in-memory
 		// to disk-backed operator occurs. It should only be set in tests.
 		SpillingCallbackFn func()
@@ -90,21 +102,16 @@ type NewColOperatorArgs struct {
 // NewColOperatorResult is a helper struct that encompasses all of the return
 // values of NewColOperator call.
 type NewColOperatorResult struct {
-	Op          colexecop.Operator
-	KVReader    colexecop.KVReader
-	ColumnTypes []*types.T
-	// MetadataSources are all sources of the metadata that are present in the
-	// tree rooted in Op that the caller must drain. These can be
-	// colfetcher.ColBatchScan or colexec.Columnarizer operators that were
-	// created during NewColOperator call, but it also might include all of the
-	// sources from NewColOperatorArgs.MetadataSources if no metadata draining
-	// component was created.
-	MetadataSources colexecop.MetadataSources
-	// ToClose is a slice of components that need to be Closed.
-	ToClose     []colexecop.Closer
-	OpMonitors  []*mon.BytesMonitor
-	OpAccounts  []*mon.BoundAccount
-	Releasables []execinfra.Releasable
+	OpWithMetaInfo
+	KVReader colexecop.KVReader
+	// Columnarizer is the root colexec.Columnarizer, if needed, that is hidden
+	// behind the stats collector interface. We need to track it separately from
+	// all other stats collectors since it requires special handling.
+	Columnarizer colexecop.VectorizedStatsCollector
+	ColumnTypes  []*types.T
+	OpMonitors   []*mon.BytesMonitor
+	OpAccounts   []*mon.BoundAccount
+	Releasables  []execinfra.Releasable
 }
 
 var _ execinfra.Releasable = &NewColOperatorResult{}
@@ -141,12 +148,14 @@ func (r *NewColOperatorResult) Release() {
 		releasable.Release()
 	}
 	*r = NewColOperatorResult{
-		ColumnTypes:     r.ColumnTypes[:0],
-		MetadataSources: r.MetadataSources[:0],
-		ToClose:         r.ToClose[:0],
-		OpMonitors:      r.OpMonitors[:0],
-		OpAccounts:      r.OpAccounts[:0],
-		Releasables:     r.Releasables[:0],
+		OpWithMetaInfo: OpWithMetaInfo{
+			MetadataSources: r.OpWithMetaInfo.MetadataSources[:0],
+			ToClose:         r.OpWithMetaInfo.ToClose[:0],
+		},
+		ColumnTypes: r.ColumnTypes[:0],
+		OpMonitors:  r.OpMonitors[:0],
+		OpAccounts:  r.OpAccounts[:0],
+		Releasables: r.Releasables[:0],
 	}
 	newColOperatorResultPool.Put(r)
 }

pkg/sql/colexec/colexecbase/ordinality_test.go

@@ -115,12 +115,12 @@ func createTestOrdinalityOperator(
 	}
 	args := &colexecargs.NewColOperatorArgs{
 		Spec:                spec,
-		Inputs:              []colexecop.Operator{input},
+		Inputs:              []colexecargs.OpWithMetaInfo{{Root: input}},
 		StreamingMemAccount: testMemAcc,
 	}
 	result, err := colexecargs.TestNewColOperator(ctx, flowCtx, args)
 	if err != nil {
 		return nil, err
 	}
-	return result.Op, nil
+	return result.Root, nil
 }

pkg/sql/colexec/colexecbase/simple_project_test.go

@@ -15,6 +15,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec"
+	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecargs"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexectestutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
@@ -116,9 +117,9 @@ func TestSimpleProjectOpWithUnorderedSynchronizer(t *testing.T) {
 	colexectestutils.RunTestsWithoutAllNullsInjection(t, testAllocator, inputTuples, [][]*types.T{inputTypes, inputTypes}, expected, colexectestutils.UnorderedVerifier,
 		func(inputs []colexecop.Operator) (colexecop.Operator, error) {
 			var input colexecop.Operator
-			parallelUnorderedSynchronizerInputs := make([]colexec.SynchronizerInput, len(inputs))
+			parallelUnorderedSynchronizerInputs := make([]colexecargs.OpWithMetaInfo, len(inputs))
 			for i := range parallelUnorderedSynchronizerInputs {
-				parallelUnorderedSynchronizerInputs[i].Op = inputs[i]
+				parallelUnorderedSynchronizerInputs[i].Root = inputs[i]
 			}
 			input = colexec.NewParallelUnorderedSynchronizer(parallelUnorderedSynchronizerInputs, &wg)
 			input = colexecbase.NewSimpleProjectOp(input, len(inputTypes), []uint32{0})

pkg/sql/colexec/colexectestutils/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "colexectestutils",
     srcs = [
+        "args.go",
         "proj_utils.go",
         "utils.go",
     ],

pkg/sql/colexec/stats.go → pkg/sql/colexec/colexectestutils/args.go

@@ -8,15 +8,19 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
-package colexec
+package colexectestutils
 
 import (
+	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecargs"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
-	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 )
 
-// VectorizedStatsCollector is the common interface implemented by collectors.
-type VectorizedStatsCollector interface {
-	colexecop.Operator
-	GetStats() *execinfrapb.ComponentStats
+// MakeInputs is a utility function that populates a slice of
+// colexecargs.OpWithMetaInfo objects based on sources.
+func MakeInputs(sources []colexecop.Operator) []colexecargs.OpWithMetaInfo {
+	inputs := make([]colexecargs.OpWithMetaInfo, len(sources))
+	for i := range sources {
+		inputs[i].Root = sources[i]
+	}
+	return inputs
 }

pkg/sql/colexec/colexectestutils/proj_utils.go

@@ -98,7 +98,7 @@ func CreateTestProjectingOperator(
 	}
 	args := &colexecargs.NewColOperatorArgs{
 		Spec:                spec,
-		Inputs:              []colexecop.Operator{input},
+		Inputs:              []colexecargs.OpWithMetaInfo{{Root: input}},
 		StreamingMemAccount: testMemAcc,
 	}
 	if canFallbackToRowexec {
@@ -108,5 +108,5 @@ func CreateTestProjectingOperator(
 	if err != nil {
 		return nil, err
 	}
-	return result.Op, nil
+	return result.Root, nil
 }

pkg/sql/colexec/colexecwindow/window_functions_test.go

@@ -274,7 +274,7 @@ func TestWindowFunctions(t *testing.T) {
 		} {
 			log.Infof(ctx, "spillForced=%t/%s", spillForced, tc.windowerSpec.WindowFns[0].Func.String())
 			var semsToCheck []semaphore.Semaphore
-			colexectestutils.RunTests(t, testAllocator, []colexectestutils.Tuples{tc.tuples}, tc.expected, colexectestutils.UnorderedVerifier, func(inputs []colexecop.Operator) (colexecop.Operator, error) {
+			colexectestutils.RunTests(t, testAllocator, []colexectestutils.Tuples{tc.tuples}, tc.expected, colexectestutils.UnorderedVerifier, func(sources []colexecop.Operator) (colexecop.Operator, error) {
 				tc.init()
 				ct := make([]*types.T, len(tc.tuples[0]))
 				for i := range ct {
@@ -297,7 +297,7 @@ func TestWindowFunctions(t *testing.T) {
 				sem := colexecop.NewTestingSemaphore(relativeRankNumRequiredFDs)
 				args := &colexecargs.NewColOperatorArgs{
 					Spec:                spec,
-					Inputs:              inputs,
+					Inputs:              colexectestutils.MakeInputs(sources),
 					StreamingMemAccount: testMemAcc,
 					DiskQueueCfg:        queueCfg,
 					FDSemaphore:         sem,
@@ -307,7 +307,7 @@ func TestWindowFunctions(t *testing.T) {
 				result, err := colexecargs.TestNewColOperator(ctx, flowCtx, args)
 				accounts = append(accounts, result.OpAccounts...)
 				monitors = append(monitors, result.OpMonitors...)
-				return result.Op, err
+				return result.Root, err
 			})
 			for i, sem := range semsToCheck {
 				require.Equal(t, 0, sem.GetCount(), "sem still reports open FDs at index %d", i)

pkg/sql/colexec/columnarizer.go

@@ -45,6 +45,8 @@ const (
 // reading the input in chunks of size coldata.BatchSize() and converting each
 // chunk into a coldata.Batch column by column.
 type Columnarizer struct {
+	// TODO(yuzefovich): consider whether embedding ProcessorBase into the
+	// columnarizers makes sense.
 	execinfra.ProcessorBase
 	colexecop.NonExplainable
 
@@ -68,6 +70,7 @@ type Columnarizer struct {
 }
 
 var _ colexecop.Operator = &Columnarizer{}
+var _ colexecop.VectorizedStatsCollector = &Columnarizer{}
 
 // NewBufferingColumnarizer returns a new Columnarizer that will be buffering up
 // rows before emitting them as output batches.
@@ -155,7 +158,37 @@ func (c *Columnarizer) Init() {
 		c.ctx = c.StartInternalNoSpan(c.ctx)
 		c.input.Start(c.ctx)
 		c.initStatus = colexecop.OperatorInitialized
+		if execStatsHijacker, ok := c.input.(execinfra.ExecStatsForTraceHijacker); ok {
+			// The columnarizer is now responsible for propagating the execution
+			// stats of the wrapped processor.
+			//
+			// Note that this columnarizer cannot be removed from the flow
+			// because it will have a vectorized stats collector planned on top,
+			// so the optimization of wrapRowSources() in execplan.go will never
+			// trigger. We check this assumption with an assertion below in the
+			// test setting.
+			//
+			// Still, just to be safe, we delay the hijacking until Init so that
+			// in case the assumption is wrong, we still get the stats from the
+			// wrapped processor.
+			c.ExecStatsForTrace = execStatsHijacker.HijackExecStatsForTrace()
+		}
+	}
+}
+
+// GetStats is part of the colexecop.VectorizedStatsCollector interface.
+func (c *Columnarizer) GetStats() *execinfrapb.ComponentStats {
+	if c.removedFromFlow && util.CrdbTestBuild {
+		colexecerror.InternalError(errors.AssertionFailedf(
+			"unexpectedly the columnarizer was removed from the flow when stats are being collected",
+		))
+	}
+	if !c.removedFromFlow && c.ExecStatsForTrace != nil {
+		s := c.ExecStatsForTrace()
+		s.Component = c.FlowCtx.ProcessorComponentID(c.ProcessorID)
+		return s
 	}
+	return nil
 }
 
 // Next is part of the Operator interface.

pkg/sql/colexec/crossjoiner_test.go

@@ -373,7 +373,7 @@ func TestCrossJoiner(t *testing.T) {
 					spec := createSpecForHashJoiner(tc)
 					args := &colexecargs.NewColOperatorArgs{
 						Spec:                spec,
-						Inputs:              sources,
+						Inputs:              colexectestutils.MakeInputs(sources),
 						StreamingMemAccount: testMemAcc,
 						DiskQueueCfg:        queueCfg,
 						FDSemaphore:         colexecop.NewTestingSemaphore(externalHJMinPartitions),
@@ -384,7 +384,7 @@ func TestCrossJoiner(t *testing.T) {
 					}
 					accounts = append(accounts, result.OpAccounts...)
 					monitors = append(monitors, result.OpMonitors...)
-					return result.Op, nil
+					return result.Root, nil
 				})
 			}
 		}
@@ -444,7 +444,7 @@ func BenchmarkCrossJoiner(b *testing.B) {
 				args := &colexecargs.NewColOperatorArgs{
 					Spec: spec,
 					// Inputs will be set below.
-					Inputs:              []colexecop.Operator{nil, nil},
+					Inputs:              []colexecargs.OpWithMetaInfo{{}, {}},
 					StreamingMemAccount: testMemAcc,
 					DiskQueueCfg:        queueCfg,
 					FDSemaphore:         colexecop.NewTestingSemaphore(VecMaxOpenFDsLimit),
@@ -459,14 +459,14 @@ func BenchmarkCrossJoiner(b *testing.B) {
 					b.SetBytes(int64(8 * nOutputRows * (len(tc.leftOutCols) + len(tc.rightOutCols))))
 					b.ResetTimer()
 					for i := 0; i < b.N; i++ {
-						args.Inputs[0] = colexectestutils.NewChunkingBatchSource(testAllocator, sourceTypes, cols, nRows)
-						args.Inputs[1] = colexectestutils.NewChunkingBatchSource(testAllocator, sourceTypes, cols, nRows)
+						args.Inputs[0].Root = colexectestutils.NewChunkingBatchSource(testAllocator, sourceTypes, cols, nRows)
+						args.Inputs[1].Root = colexectestutils.NewChunkingBatchSource(testAllocator, sourceTypes, cols, nRows)
 						result, err := colexecargs.TestNewColOperator(ctx, flowCtx, args)
 						require.NoError(b, err)
 						accounts = append(accounts, result.OpAccounts...)
 						monitors = append(monitors, result.OpMonitors...)
 						require.NoError(b, err)
-						cj := result.Op
+						cj := result.Root
 						cj.Init()
 						for b := cj.Next(ctx); b.Length() > 0; b = cj.Next(ctx) {
 						}

pkg/sql/colexec/external_distinct_test.go

@@ -352,7 +352,7 @@ func BenchmarkExternalDistinct(b *testing.B) {
 func createExternalDistinct(
 	ctx context.Context,
 	flowCtx *execinfra.FlowCtx,
-	input []colexecop.Operator,
+	sources []colexecop.Operator,
 	typs []*types.T,
 	distinctCols []uint32,
 	outputOrdering execinfrapb.Ordering,
@@ -375,13 +375,13 @@ func createExternalDistinct(
 	}
 	args := &colexecargs.NewColOperatorArgs{
 		Spec:                spec,
-		Inputs:              input,
+		Inputs:              colexectestutils.MakeInputs(sources),
 		StreamingMemAccount: testMemAcc,
 		DiskQueueCfg:        diskQueueCfg,
 		FDSemaphore:         testingSemaphore,
 	}
 	args.TestingKnobs.SpillingCallbackFn = spillingCallbackFn
 	args.TestingKnobs.NumForcedRepartitions = numForcedRepartitions
 	result, err := colexecargs.TestNewColOperator(ctx, flowCtx, args)
-	return result.Op, result.OpAccounts, result.OpMonitors, result.ToClose, err
+	return result.Root, result.OpAccounts, result.OpMonitors, result.ToClose, err
 }