Merge #145934

145934: sql: refactor plan node control flow r=DrewKimball a=DrewKimball

#### sql: change SetZoneConfig to return type Ack

This commit changes the return type for `tree.SetZoneConfig` to `Ack`
from `RowsAffected`. This mirrors similar statements, and the row count
has been broken (always zero) ever since this statement was implemented
in the declarative schema changer, anyway. Removing the need to track
row count allows a few related functions to be simplified.

Epic: None

Release note: None

#### sql: remove unnecessary nodes and interfaces related to mutation planNodes

This commit adds two new helper structs for use in mutation nodes and
other planNodes that need to output the number of affected rows:
`mutationOutputHelper` and `rowsAffectedOutputHelper`. The former is
now embedded into the mutation planNodes apart from `deleteRangeNode`,
which embeds `rowsAffectedOutputHelper` instead.

This commit also removes `spoolNode`, `serializeNode`, and `rowCount`
node, with `mutationOutputHelper` replacing the work they used to do.
This also allowed the removal of the `batchedPlanNode` interface and
its methods.

These changes simplify some factory logic, and should simplify the
process of moving from `planNode` implementations to processors in
the future.

This change also fixes #139378, where the `spoolNode` added for a mutation
with a RETURNING clause would break some of the EXPLAIN ANALYZE stats.

Fixes #139378

Release note (bug fix): Previously, CockroachDB would omit execution
statistics in EXPLAIN ANALYZE output for mutation nodes when a RETURNING
clause was used. The bug has been present since before v21.1 and is now fixed.

#### sql: remove planNodeFastPath interface for ops that return rows affected

Previously, there were two ways that `planNode` implementations would
propagate the number of rows affected (if using that output mode):
1. Implementing the `planNodeFastPath`, which provided a method to
   directly pull the cached row count.
2. A special loop in `planNodeToRowSource`, which simply called `Next()`
   repeatedly and counted the number of iterations.

This commit takes advantage of work done in the previous commit for
operators that need to return the number of rows affected. Now, the
helper structs from the previous commit are used to track this count
internally, and then return the final count through the usual `Next()`
and `Values()` calls. This allows for the removal of special logic
in `planNodeToRowSource`, as well as the `planNodeFastPath` interface.

As with the previous commit, this one seeks to remove departures from
the usual `Start()`, `Next()`, and `Values()` control flow, to make
future refactors easier.

Epic: None

Release note: None

Co-authored-by: Drew Kimball <[email protected]>

Author: craig[bot]

pkg/cli/interactive_tests/test_missing_log_output.tcl

@@ -56,7 +56,7 @@ interrupt_and_wait
 # Disable replication so as to avoid spurious purgatory errors.
 start_server $argv
 send "$argv sql --insecure -e \"ALTER RANGE default CONFIGURE ZONE USING num_replicas = 1\"\r"
-eexpect {CONFIGURE ZONE [0-9]}
+eexpect "CONFIGURE ZONE"
 eexpect ":/# "
 stop_server $argv
 

pkg/sql/BUILD.bazel

@@ -171,9 +171,9 @@ go_library(
         "pg_extension.go",
         "pg_metadata_diff.go",
         "plan.go",
-        "plan_batch.go",
         "plan_columns.go",
         "plan_names.go",
+        "plan_node_output_helper.go",
         "plan_node_to_row_source.go",
         "plan_opt.go",
         "plan_ordering.go",
@@ -248,7 +248,6 @@ go_library(
         "show_zone_config.go",
         "sort.go",
         "split.go",
-        "spool.go",
         "sql_activity_update_job.go",
         "sql_cursor.go",
         "statement.go",
@@ -722,6 +721,7 @@ go_test(
         "pg_metadata_test.go",
         "pg_oid_test.go",
         "pgwire_internal_test.go",
+        "plan_node_output_helper_test.go",
         "plan_opt_test.go",
         "privileged_accessor_test.go",
         "region_util_test.go",
@@ -893,6 +893,7 @@ go_test(
         "//pkg/sql/randgen",
         "//pkg/sql/regions",
         "//pkg/sql/row",
+        "//pkg/sql/rowcontainer",
         "//pkg/sql/rowenc",
         "//pkg/sql/rowenc/keyside",
         "//pkg/sql/rowenc/rowencpb",

pkg/sql/cancel_queries.go

@@ -19,14 +19,23 @@ import (
 
 type cancelQueriesNode struct {
 	singleInputPlanNode
+	rowsAffectedOutputHelper
 	ifExists bool
 }
 
-func (n *cancelQueriesNode) startExec(runParams) error {
-	return nil
+// startExec implements the planNode interface.
+func (n *cancelQueriesNode) startExec(params runParams) error {
+	// Execute the node to completion, keeping track of the affected row count.
+	for {
+		ok, err := n.cancelQuery(params)
+		if !ok || err != nil {
+			return err
+		}
+		n.incAffectedRows()
+	}
 }
 
-func (n *cancelQueriesNode) Next(params runParams) (bool, error) {
+func (n *cancelQueriesNode) cancelQuery(params runParams) (bool, error) {
 	// TODO(knz): instead of performing the cancels sequentially,
 	// accumulate all the query IDs and then send batches to each of the
 	// nodes.
@@ -71,7 +80,15 @@ func (n *cancelQueriesNode) Next(params runParams) (bool, error) {
 	return true, nil
 }
 
-func (*cancelQueriesNode) Values() tree.Datums { return nil }
+// Next implements the planNode interface.
+func (n *cancelQueriesNode) Next(_ runParams) (bool, error) {
+	return n.next(), nil
+}
+
+// Values implements the planNode interface.
+func (n *cancelQueriesNode) Values() tree.Datums {
+	return n.values()
+}
 
 func (n *cancelQueriesNode) Close(ctx context.Context) {
 	n.input.Close(ctx)

pkg/sql/cancel_sessions.go

@@ -19,14 +19,23 @@ import (
 
 type cancelSessionsNode struct {
 	singleInputPlanNode
+	rowsAffectedOutputHelper
 	ifExists bool
 }
 
-func (n *cancelSessionsNode) startExec(runParams) error {
-	return nil
+// startExec implements the planNode interface.
+func (n *cancelSessionsNode) startExec(params runParams) error {
+	// Execute the node to completion, keeping track of the affected row count.
+	for {
+		ok, err := n.cancelSession(params)
+		if !ok || err != nil {
+			return err
+		}
+		n.incAffectedRows()
+	}
 }
 
-func (n *cancelSessionsNode) Next(params runParams) (bool, error) {
+func (n *cancelSessionsNode) cancelSession(params runParams) (bool, error) {
 	// TODO(knz): instead of performing the cancels sequentially,
 	// accumulate all the query IDs and then send batches to each of the
 	// nodes.
@@ -71,7 +80,15 @@ func (n *cancelSessionsNode) Next(params runParams) (bool, error) {
 	return true, nil
 }
 
-func (*cancelSessionsNode) Values() tree.Datums { return nil }
+// Next implements the planNode interface.
+func (n *cancelSessionsNode) Next(_ runParams) (bool, error) {
+	return n.next(), nil
+}
+
+// Values implements the planNode interface.
+func (n *cancelSessionsNode) Values() tree.Datums {
+	return n.values()
+}
 
 func (n *cancelSessionsNode) Close(ctx context.Context) {
 	n.input.Close(ctx)

pkg/sql/colexec/columnarizer.go

@@ -211,18 +211,18 @@ func (c *Columnarizer) GetStats() *execinfrapb.ComponentStats {
 	return s
 }
 
-// IsFastPathNode returns true if the provided RowSource is the
-// planNodeToRowSource wrapper with "fast-path" enabled. The logic is injected
-// in the sql package to avoid an import cycle.
-var IsFastPathNode func(source execinfra.RowSource) bool
+// IsRowsAffectedNode returns true if the provided RowSource is the
+// planNodeToRowSource wrapper returning the number of affected rows. The logic
+// is injected from the sql package to avoid an import cycle.
+var IsRowsAffectedNode func(source execinfra.RowSource) bool
 
-// IsColumnarizerAroundFastPathNode returns true if the provided Operator is a
-// Columnarizer that has "fast-path" enabled planNodeToRowSource wrapper as the
-// input.
-func IsColumnarizerAroundFastPathNode(o colexecop.Operator) bool {
+// IsColumnarizerAroundRowsAffectedNode returns true if the provided Operator is
+// a Columnarizer that has a planNodeToRowSource wrapper in "rows affected" mode
+// as the input.
+func IsColumnarizerAroundRowsAffectedNode(o colexecop.Operator) bool {
 	o = MaybeUnwrapInvariantsChecker(o)
 	c, ok := o.(*Columnarizer)
-	return ok && IsFastPathNode(c.input)
+	return ok && IsRowsAffectedNode(c.input)
 }
 
 // Next is part of the colexecop.Operator interface.

pkg/sql/colflow/stats.go

@@ -57,9 +57,9 @@ type batchInfoCollector struct {
 	// batch is the last batch returned by the wrapped operator.
 	batch coldata.Batch
 
-	// rowCountFastPath is set to indicate that the input is expected to produce
+	// rowsAffectedMode is set to indicate that the input is expected to produce
 	// a single batch with a single column with the row count value.
-	rowCountFastPath bool
+	rowsAffectedMode bool
 
 	// stopwatch keeps track of the amount of time the wrapped operator spent
 	// doing work. Note that this will include all of the time that the operator's
@@ -86,7 +86,7 @@ func makeBatchInfoCollector(
 	return batchInfoCollector{
 		OneInputNode:         colexecop.NewOneInputNode(op),
 		componentID:          id,
-		rowCountFastPath:     colexec.IsColumnarizerAroundFastPathNode(op),
+		rowsAffectedMode:     colexec.IsColumnarizerAroundRowsAffectedNode(op),
 		stopwatch:            inputWatch,
 		childStatsCollectors: childStatsCollectors,
 	}
@@ -134,12 +134,12 @@ func (bic *batchInfoCollector) Next() coldata.Batch {
 		bic.mu.Lock()
 		defer bic.mu.Unlock()
 		bic.mu.numBatches++
-		if bic.rowCountFastPath {
+		if bic.rowsAffectedMode {
 			// We have a special case where the batch has exactly one column
 			// with exactly one row in which we have the row count.
 			if buildutil.CrdbTestBuild {
 				if bic.mu.numBatches != 1 {
-					colexecerror.InternalError(errors.AssertionFailedf("saw second batch in fast path:\n%s", bic.batch))
+					colexecerror.InternalError(errors.AssertionFailedf("saw second batch in rows affected mode:\n%s", bic.batch))
 				}
 				if bic.batch.Width() != 1 {
 					colexecerror.InternalError(errors.AssertionFailedf("batch width is not 1:\n%s", bic.batch))

pkg/sql/control_jobs.go

@@ -20,8 +20,8 @@ import (
 
 type controlJobsNode struct {
 	singleInputPlanNode
+	rowsAffectedOutputHelper
 	desiredStatus jobs.State
-	numRows       int
 	reason        string
 }
 
@@ -31,11 +31,7 @@ var jobCommandToDesiredStatus = map[tree.JobCommand]jobs.State{
 	tree.PauseJob:  jobs.StatePaused,
 }
 
-// FastPathResults implements the planNodeFastPath interface.
-func (n *controlJobsNode) FastPathResults() (int, bool) {
-	return n.numRows, true
-}
-
+// startExec implements the planNode interface.
 func (n *controlJobsNode) startExec(params runParams) error {
 	if n.desiredStatus != jobs.StatePaused && len(n.reason) > 0 {
 		return errors.AssertionFailedf("status %v is not %v and thus does not support a reason %v",
@@ -86,7 +82,7 @@ func (n *controlJobsNode) startExec(params runParams) error {
 			return err
 		}
 
-		n.numRows++
+		n.incAffectedRows()
 	}
 	switch n.desiredStatus {
 	case jobs.StatePaused:
@@ -99,9 +95,15 @@ func (n *controlJobsNode) startExec(params runParams) error {
 	return nil
 }
 
-func (*controlJobsNode) Next(runParams) (bool, error) { return false, nil }
+// Next implements the planNode interface.
+func (n *controlJobsNode) Next(_ runParams) (bool, error) {
+	return n.next(), nil
+}
 
-func (*controlJobsNode) Values() tree.Datums { return nil }
+// Values implements the planNode interface.
+func (n *controlJobsNode) Values() tree.Datums {
+	return n.values()
+}
 
 func (n *controlJobsNode) Close(ctx context.Context) {
 	n.input.Close(ctx)

pkg/sql/control_schedules.go

@@ -26,8 +26,8 @@ import (
 
 type controlSchedulesNode struct {
 	singleInputPlanNode
+	rowsAffectedOutputHelper
 	command tree.ScheduleCommand
-	numRows int
 }
 
 func collectTelemetry(command tree.ScheduleCommand) {
@@ -41,11 +41,6 @@ func collectTelemetry(command tree.ScheduleCommand) {
 	}
 }
 
-// FastPathResults implements the planNodeFastPath interface.
-func (n *controlSchedulesNode) FastPathResults() (int, bool) {
-	return n.numRows, true
-}
-
 // JobSchedulerEnv returns JobSchedulerEnv.
 func JobSchedulerEnv(knobs *jobs.TestingKnobs) scheduledjobs.JobSchedulerEnv {
 	if knobs != nil && knobs.JobSchedulerEnv != nil {
@@ -175,7 +170,7 @@ func (n *controlSchedulesNode) startExec(params runParams) error {
 				if err != nil {
 					return errors.Wrap(err, "failed to run OnDrop")
 				}
-				n.numRows += additionalDroppedSchedules
+				n.addAffectedRows(additionalDroppedSchedules)
 			}
 			err = DeleteSchedule(
 				params.ctx, params.ExecCfg(), params.p.InternalSQLTxn(),
@@ -189,17 +184,21 @@ func (n *controlSchedulesNode) startExec(params runParams) error {
 		if err != nil {
 			return err
 		}
-		n.numRows++
+		n.incAffectedRows()
 	}
 
 	return nil
 }
 
-// Next implements planNode interface.
-func (*controlSchedulesNode) Next(runParams) (bool, error) { return false, nil }
+// Next implements the planNode interface.
+func (n *controlSchedulesNode) Next(_ runParams) (bool, error) {
+	return n.next(), nil
+}
 
-// Values implements planNode interface.
-func (*controlSchedulesNode) Values() tree.Datums { return nil }
+// Values implements the planNode interface.
+func (n *controlSchedulesNode) Values() tree.Datums {
+	return n.values()
+}
 
 // Close implements planNode interface.
 func (n *controlSchedulesNode) Close(ctx context.Context) {

pkg/sql/create_table.go

@@ -570,7 +570,6 @@ func (n *createTableNode) startExec(params runParams) error {
 			ti := tableInserterPool.Get().(*tableInserter)
 			*ti = tableInserter{ri: ri}
 			defer func() {
-				ti.close(params.ctx)
 				*ti = tableInserter{}
 				tableInserterPool.Put(ti)
 			}()

pkg/sql/database_region_change_finalizer.go

@@ -115,7 +115,7 @@ func (r *databaseRegionChangeFinalizer) preDrop(ctx context.Context, txn descs.T
 		return err
 	}
 	for _, update := range zoneConfigUpdates {
-		if _, err := writeZoneConfigUpdate(
+		if err = writeZoneConfigUpdate(
 			ctx, txn,
 			r.localPlanner.ExtendedEvalContext().Tracing.KVTracingEnabled(),
 			update,