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.

Author: DrewKimball

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

@@ -36,12 +36,12 @@ var _ mutationPlanNode = &deleteNode{}
 
 // deleteRun contains the run-time state of deleteNode during local execution.
 type deleteRun struct {
-	td         tableDeleter
-	rowsNeeded bool
+	mutationOutputHelper
+	td tableDeleter
 
-	// done informs a new call to BatchedNext() that the previous call
-	// to BatchedNext() has completed the work already.
-	done bool
+	// rowsNeeded is set to true if the mutation operator needs to return the rows
+	// that were affected by the mutation.
+	rowsNeeded bool
 
 	// resultRowBuffer is used to prepare a result row for accumulation
 	// into the row container above, when rowsNeeded is set.
@@ -76,7 +76,7 @@ func (r *deleteRun) init(params runParams, columns colinfo.ResultColumns) {
 		return
 	}
 
-	r.td.rows = rowcontainer.NewRowContainer(
+	r.rows = rowcontainer.NewRowContainer(
 		params.p.Mon().MakeBoundAccount(),
 		colinfo.ColTypeInfoFromResCols(columns),
 	)
@@ -92,31 +92,34 @@ func (d *deleteNode) startExec(params runParams) error {
 
 	d.run.init(params, d.columns)
 
-	return d.run.td.init(params.ctx, params.p.txn, params.EvalContext())
-}
+	if err := d.run.td.init(params.ctx, params.p.txn, params.EvalContext()); err != nil {
+		return err
+	}
 
-// Next is required because batchedPlanNode inherits from planNode, but
-// batchedPlanNode doesn't really provide it. See the explanatory comments
-// in plan_batch.go.
-func (d *deleteNode) Next(params runParams) (bool, error) { panic("not valid") }
+	// Run the mutation to completion.
+	for {
+		lastBatch, err := d.processBatch(params)
+		if err != nil || lastBatch {
+			return err
+		}
+	}
+}
 
-// Values is required because batchedPlanNode inherits from planNode, but
-// batchedPlanNode doesn't really provide it. See the explanatory comments
-// in plan_batch.go.
-func (d *deleteNode) Values() tree.Datums { panic("not valid") }
+// Next implements the planNode interface.
+func (d *deleteNode) Next(_ runParams) (bool, error) {
+	return d.run.next(), nil
+}
 
-// BatchedNext implements the batchedPlanNode interface.
-func (d *deleteNode) BatchedNext(params runParams) (bool, error) {
-	if d.run.done {
-		return false, nil
-	}
+// Values implements the planNode interface.
+func (d *deleteNode) Values() tree.Datums {
+	return d.run.values()
+}
 
-	// Advance one batch. First, clear the last batch.
-	d.run.td.clearLastBatch(params.ctx)
-	// Now consume/accumulate the rows for this batch.
-	lastBatch := false
+func (d *deleteNode) processBatch(params runParams) (lastBatch bool, err error) {
+	// Consume/accumulate the rows for this batch.
+	lastBatch = false
 	for {
-		if err := params.p.cancelChecker.Check(); err != nil {
+		if err = params.p.cancelChecker.Check(); err != nil {
 			return false, err
 		}
 
@@ -131,7 +134,7 @@ func (d *deleteNode) BatchedNext(params runParams) (bool, error) {
 
 		// Process the deletion of the current input row,
 		// potentially accumulating the result row for later.
-		if err := d.run.processSourceRow(params, d.input.Values()); err != nil {
+		if err = d.run.processSourceRow(params, d.input.Values()); err != nil {
 			return false, err
 		}
 
@@ -146,24 +149,21 @@ func (d *deleteNode) BatchedNext(params runParams) (bool, error) {
 		if !lastBatch {
 			// We only run/commit the batch if there were some rows processed
 			// in this batch.
-			if err := d.run.td.flushAndStartNewBatch(params.ctx); err != nil {
+			if err = d.run.td.flushAndStartNewBatch(params.ctx); err != nil {
 				return false, err
 			}
 		}
 	}
 
 	if lastBatch {
 		d.run.td.setRowsWrittenLimit(params.extendedEvalCtx.SessionData())
-		if err := d.run.td.finalize(params.ctx); err != nil {
+		if err = d.run.td.finalize(params.ctx); err != nil {
 			return false, err
 		}
-		// Remember we're done for the next call to BatchedNext().
-		d.run.done = true
 		// Possibly initiate a run of CREATE STATISTICS.
-		params.ExecCfg().StatsRefresher.NotifyMutation(d.run.td.tableDesc(), int(d.run.td.rowsWritten))
+		params.ExecCfg().StatsRefresher.NotifyMutation(d.run.td.tableDesc(), int(d.run.rowsAffected()))
 	}
-
-	return d.run.td.lastBatchSize > 0, nil
+	return lastBatch, nil
 }
 
 // processSourceRow processes one row from the source for deletion and, if
@@ -199,64 +199,60 @@ func (r *deleteRun) processSourceRow(params runParams, sourceVals tree.Datums) e
 	); err != nil {
 		return err
 	}
+	r.onModifiedRow()
+	if !r.rowsNeeded {
+		return nil
+	}
 
-	// If result rows need to be accumulated, do it.
-	if r.td.rows != nil {
-		// The new values can include all columns, so the values may contain
-		// additional columns for every newly dropped column not visible. We do not
-		// want them to be available for RETURNING.
-		//
-		// r.rows.NumCols() is guaranteed to only contain the requested
-		// public columns.
-		largestRetIdx := -1
-		for i := range r.rowIdxToRetIdx {
-			retIdx := r.rowIdxToRetIdx[i]
-			if retIdx >= 0 {
-				if retIdx >= largestRetIdx {
-					largestRetIdx = retIdx
-				}
-				r.resultRowBuffer[retIdx] = deleteVals[i]
+	// Result rows must be accumulated.
+	//
+	// The new values can include all columns, so the values may contain
+	// additional columns for every newly dropped column not visible. We do not
+	// want them to be available for RETURNING.
+	//
+	// r.rows.NumCols() is guaranteed to only contain the requested
+	// public columns.
+	largestRetIdx := -1
+	for i := range r.rowIdxToRetIdx {
+		retIdx := r.rowIdxToRetIdx[i]
+		if retIdx >= 0 {
+			if retIdx >= largestRetIdx {
+				largestRetIdx = retIdx
 			}
+			r.resultRowBuffer[retIdx] = deleteVals[i]
 		}
+	}
 
-		// At this point we've extracted all the RETURNING values that are part
-		// of the target table. We must now extract the columns in the RETURNING
-		// clause that refer to other tables (from the USING clause of the delete).
-		if r.numPassthrough > 0 {
-			passthroughBegin := len(r.td.rd.FetchCols)
-			passthroughEnd := passthroughBegin + r.numPassthrough
-			passthroughValues := deleteVals[passthroughBegin:passthroughEnd]
-
-			for i := 0; i < r.numPassthrough; i++ {
-				largestRetIdx++
-				r.resultRowBuffer[largestRetIdx] = passthroughValues[i]
-			}
-
+	// At this point we've extracted all the RETURNING values that are part
+	// of the target table. We must now extract the columns in the RETURNING
+	// clause that refer to other tables (from the USING clause of the delete).
+	if r.numPassthrough > 0 {
+		passthroughBegin := len(r.td.rd.FetchCols)
+		passthroughEnd := passthroughBegin + r.numPassthrough
+		passthroughValues := deleteVals[passthroughBegin:passthroughEnd]
+
+		for i := 0; i < r.numPassthrough; i++ {
+			largestRetIdx++
+			r.resultRowBuffer[largestRetIdx] = passthroughValues[i]
 		}
 
-		if _, err := r.td.rows.AddRow(params.ctx, r.resultRowBuffer); err != nil {
-			return err
-		}
 	}
-
-	return nil
+	return r.addRow(params.ctx, r.resultRowBuffer)
 }
 
-// BatchedCount implements the batchedPlanNode interface.
-func (d *deleteNode) BatchedCount() int { return d.run.td.lastBatchSize }
-
-// BatchedValues implements the batchedPlanNode interface.
-func (d *deleteNode) BatchedValues(rowIdx int) tree.Datums { return d.run.td.rows.At(rowIdx) }
-
 func (d *deleteNode) Close(ctx context.Context) {
 	d.input.Close(ctx)
-	d.run.td.close(ctx)
+	d.run.close(ctx)
 	*d = deleteNode{}
 	deleteNodePool.Put(d)
 }
 
 func (d *deleteNode) rowsWritten() int64 {
-	return d.run.td.rowsWritten
+	return d.run.rowsAffected()
+}
+
+func (d *deleteNode) returnsRowsAffected() bool {
+	return !d.run.rowsNeeded
 }
 
 func (d *deleteNode) enableAutoCommit() {

pkg/sql/delete_range.go

@@ -31,6 +31,7 @@ import (
 // be deleted, it'll enable autoCommit for delete range.
 type deleteRangeNode struct {
 	zeroInputPlanNode
+	rowsAffectedOutputHelper
 	// spans are the spans to delete.
 	spans roachpb.Spans
 	// desc is the table descriptor the delete is operating on.
@@ -45,42 +46,22 @@ type deleteRangeNode struct {
 	// batches and will just send one big delete with a commit statement attached.
 	autoCommitEnabled bool
 
-	// rowCount will be set to the count of rows deleted.
-	rowCount int
-
 	// curRowPrefix is the prefix for all KVs (i.e. for all column families) of
 	// the SQL row that increased rowCount last. It is maintained across
 	// different BatchRequests in order to not double count the same SQL row.
 	curRowPrefix []byte
 }
 
 var _ planNode = &deleteRangeNode{}
-var _ planNodeFastPath = &deleteRangeNode{}
-var _ batchedPlanNode = &deleteRangeNode{}
 var _ mutationPlanNode = &deleteRangeNode{}
 
-// BatchedNext implements the batchedPlanNode interface.
-func (d *deleteRangeNode) BatchedNext(params runParams) (bool, error) {
-	return false, nil
-}
-
-// BatchedCount implements the batchedPlanNode interface.
-func (d *deleteRangeNode) BatchedCount() int {
-	return d.rowCount
-}
-
-// BatchedValues implements the batchedPlanNode interface.
-func (d *deleteRangeNode) BatchedValues(rowIdx int) tree.Datums {
-	panic("invalid")
-}
-
-// FastPathResults implements the planNodeFastPath interface.
-func (d *deleteRangeNode) FastPathResults() (int, bool) {
-	return d.rowCount, true
+func (d *deleteRangeNode) rowsWritten() int64 {
+	return d.rowsAffected()
 }
 
-func (d *deleteRangeNode) rowsWritten() int64 {
-	return int64(d.rowCount)
+func (d *deleteRangeNode) returnsRowsAffected() bool {
+	// DeleteRange always returns the number of rows deleted.
+	return true
 }
 
 // startExec implements the planNode interface.
@@ -228,7 +209,7 @@ func (d *deleteRangeNode) processResults(
 			k := keyBytes[:len(keyBytes)-len(after)]
 			if !bytes.Equal(k, d.curRowPrefix) {
 				d.curRowPrefix = k
-				d.rowCount++
+				d.incAffectedRows()
 			}
 		}
 		if r.ResumeSpan != nil && r.ResumeSpan.Valid() {
@@ -239,15 +220,13 @@ func (d *deleteRangeNode) processResults(
 }
 
 // Next implements the planNode interface.
-func (*deleteRangeNode) Next(params runParams) (bool, error) {
-	// TODO(radu): this shouldn't be used, but it gets called when a cascade uses
-	// delete-range. Investigate this.
-	return false, nil
+func (d *deleteRangeNode) Next(params runParams) (bool, error) {
+	return d.next(), nil
 }
 
 // Values implements the planNode interface.
-func (*deleteRangeNode) Values() tree.Datums {
-	panic("invalid")
+func (d *deleteRangeNode) Values() tree.Datums {
+	return d.values()
 }
 
 // Close implements the planNode interface.