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

Author: DrewKimball

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/distsql_physical_planner.go

@@ -4285,8 +4285,6 @@ func (dsp *DistSQLPlanner) createPhysPlanForPlanNode(
 func (dsp *DistSQLPlanner) wrapPlan(
 	ctx context.Context, planCtx *PlanningCtx, n planNode, allowPartialDistribution bool,
 ) (*PhysicalPlan, error) {
-	useFastPath := planCtx.planDepth == 1 && planCtx.stmtType == tree.RowsAffected
-
 	// First, we search the planNode tree we're trying to wrap for the first
 	// DistSQL-enabled planNode in the tree. If we find one, we ask the planner to
 	// continue the DistSQL planning recursion on that planNode.
@@ -4341,20 +4339,19 @@ func (dsp *DistSQLPlanner) wrapPlan(
 
 	// Copy the evalCtx.
 	evalCtx := *planCtx.ExtendedEvalCtx
-	// We permit the planNodeToRowSource to trigger the wrapped planNode's fast
-	// path if its the very first node in the flow, and if the statement type we're
-	// expecting is in fact RowsAffected. RowsAffected statements return a single
-	// row with the number of rows affected by the statement, and are the only
-	// types of statement where it's valid to invoke a plan's fast path.
 	wrapper := newPlanNodeToRowSource(
 		n,
 		runParams{
 			extendedEvalCtx: &evalCtx,
 			p:               planCtx.planner,
 		},
-		useFastPath,
 		firstNotWrapped,
 	)
+	if !wrapper.rowsAffected && planCtx.planDepth == 1 && planCtx.stmtType == tree.RowsAffected {
+		// Return an error if the receiver expects to get the number of rows
+		// affected, but the planNode returns something else.
+		return nil, errors.AssertionFailedf("planNode %T should return rows affected", n)
+	}
 
 	localProcIdx := p.AddLocalProcessor(wrapper)
 	var input []execinfrapb.InputSyncSpec

pkg/sql/opaque.go

@@ -238,7 +238,7 @@ func planOpaque(ctx context.Context, p *planner, stmt tree.Statement) (planNode,
 	case *tree.Inspect:
 		return p.Inspect(ctx, n)
 	case *tree.MoveCursor:
-		return p.FetchCursor(ctx, &n.CursorStmt)
+		return p.MoveCursor(ctx, &n.CursorStmt)
 	case *tree.ReassignOwnedBy:
 		return p.ReassignOwnedBy(ctx, n)
 	case *tree.RefreshMaterializedView:

pkg/sql/opt/exec/execbuilder/testdata/unique

@@ -2679,7 +2679,7 @@ quality of service: regular
 │       ├── • update
 │       │   │ sql nodes: <hidden>
 │       │   │ regions: <hidden>
-│       │   │ actual row count: 1
+│       │   │ actual row count: 2
 │       │   │ execution time: 0µs
 │       │   │ table: uniq_fk_child
 │       │   │ set: b, c

pkg/sql/plan.go

@@ -75,8 +75,7 @@ type planNode interface {
 	// of results each time that Next() returns true.
 	//
 	// Available after startPlan(). It is illegal to call Next() after it returns
-	// false. It is legal to call Next() even if the node implements
-	// planNodeFastPath and the FastPathResults() method returns true.
+	// false.
 	Next(params runParams) (bool, error)
 
 	// Values returns the values at the current row. The result is only valid
@@ -156,17 +155,6 @@ type mutationPlanNode interface {
 	returnsRowsAffected() bool
 }
 
-// planNodeFastPath is implemented by nodes that can perform all their
-// work during startPlan(), possibly affecting even multiple rows. For
-// example, DELETE can do this.
-type planNodeFastPath interface {
-	// FastPathResults returns the affected row count and true if the
-	// node has no result set and has already executed when startPlan() completes.
-	// Note that Next() must still be valid even if this method returns
-	// true, although it may have nothing left to do.
-	FastPathResults() (int, bool)
-}
-
 // planNodeReadingOwnWrites can be implemented by planNodes which do
 // not use the standard SQL principle of reading at the snapshot
 // established at the start of the transaction. It requests that
@@ -271,9 +259,6 @@ var _ planNode = &virtualTableNode{}
 var _ planNode = &windowNode{}
 var _ planNode = &zeroNode{}
 
-var _ planNodeFastPath = &controlJobsNode{}
-var _ planNodeFastPath = &controlSchedulesNode{}
-
 var _ planNodeReadingOwnWrites = &alterIndexNode{}
 var _ planNodeReadingOwnWrites = &alterSchemaNode{}
 var _ planNodeReadingOwnWrites = &alterSequenceNode{}