Merge #144054 #144071

144054: jsonpath: optimize executeAnyItem by iterating over JSON containers r=normanchenn a=normanchenn

This commit replaces the usage of `AllPathsWithDepth` with direct
iteration over JSON arrays and objects in `executeAnyItem`. This avoids
unnecessary JSON decoding of the entire object structure
(`jsonEncoded.decode()`) when only direct children are needed.

Preliminary testing shows this change reduces execution time for queries
with any array/object unwrapping by approximately 50%.

Epic: None
Release note: None

144071: sql: refactor updateRun, upsertRun, and deleteRun r=yuzefovich a=michae2

**sql: bring back resultRowBuffer in updateRun and deleteRun**

For some reason, updateRun and deleteRun differ from insertRun in that
they were creating result slices for RETURNING in processSourceRow,
instead of creating a resultRowBuffer in startExec. This commit does a
small refactor to move allocation of resultRowBuffer to startExec.

This refactor will make it easier for the new UPDATE and DELETE fast
path planNodes to share updateRun and deleteRun with the existing
updateNode and deleteNode.

Epic: None

Release note: None
    
---

**sql: move processSourceRow to updateRun, upsertRun, and deleteRun**

Refactor processSourceRow so that it's a method of updateRun, upsertRun,
and deleteRun, rather than a method of updateNode, upsertNode, and
deleteNode, respectively. This matches insertRun.

This refactor will make it easier for the new UPDATE and DELETE fast
path planNodes to share updateRun and deleteRun with the existing
updateNode and deleteNode. (There's no new fast path in the works for
UPSERT, at least not yet, but upsertRun is changed for completeness.)

Epic: None

Release note: None

Co-authored-by: Norman Chen <[email protected]>
Co-authored-by: Michael Erickson <[email protected]>

Author: 3 people

pkg/sql/delete.go

@@ -32,6 +32,8 @@ type deleteNode struct {
 	run deleteRun
 }
 
+var _ mutationPlanNode = &deleteNode{}
+
 // deleteRun contains the run-time state of deleteNode during local execution.
 type deleteRun struct {
 	td         tableDeleter
@@ -41,6 +43,10 @@ type deleteRun struct {
 	// to BatchedNext() has completed the work already.
 	done bool
 
+	// resultRowBuffer is used to prepare a result row for accumulation
+	// into the row container above, when rowsNeeded is set.
+	resultRowBuffer tree.Datums
+
 	// traceKV caches the current KV tracing flag.
 	traceKV bool
 
@@ -57,17 +63,26 @@ type deleteRun struct {
 	numPassthrough int
 }
 
-var _ mutationPlanNode = &deleteNode{}
+func (r *deleteRun) initRowContainer(params runParams, columns colinfo.ResultColumns) {
+	if !r.rowsNeeded {
+		return
+	}
+	r.td.rows = rowcontainer.NewRowContainer(
+		params.p.Mon().MakeBoundAccount(),
+		colinfo.ColTypeInfoFromResCols(columns),
+	)
+	r.resultRowBuffer = make([]tree.Datum, len(columns))
+	for i := range r.resultRowBuffer {
+		r.resultRowBuffer[i] = tree.DNull
+	}
+}
 
 func (d *deleteNode) startExec(params runParams) error {
 	// cache traceKV during execution, to avoid re-evaluating it for every row.
 	d.run.traceKV = params.p.ExtendedEvalContext().Tracing.KVTracingEnabled()
 
-	if d.run.rowsNeeded {
-		d.run.td.rows = rowcontainer.NewRowContainer(
-			params.p.Mon().MakeBoundAccount(),
-			colinfo.ColTypeInfoFromResCols(d.columns))
-	}
+	d.run.initRowContainer(params, d.columns)
+
 	return d.run.td.init(params.ctx, params.p.txn, params.EvalContext())
 }
 
@@ -107,7 +122,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.processSourceRow(params, d.input.Values()); err != nil {
+		if err := d.run.processSourceRow(params, d.input.Values()); err != nil {
 			return false, err
 		}
 
@@ -145,18 +160,18 @@ func (d *deleteNode) BatchedNext(params runParams) (bool, error) {
 
 // processSourceRow processes one row from the source for deletion and, if
 // result rows are needed, saves it in the result row container
-func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums) error {
+func (r *deleteRun) processSourceRow(params runParams, sourceVals tree.Datums) error {
 	// Remove extra columns for partial index predicate values and AFTER triggers.
-	deleteVals := sourceVals[:len(d.run.td.rd.FetchCols)+d.run.numPassthrough]
+	deleteVals := sourceVals[:len(r.td.rd.FetchCols)+r.numPassthrough]
 	sourceVals = sourceVals[len(deleteVals):]
 
 	// Create a set of partial index IDs to not delete from. Indexes should not
 	// be deleted from when they are partial indexes and the row does not
 	// satisfy the predicate and therefore do not exist in the partial index.
 	// This set is passed as a argument to tableDeleter.row below.
 	var pm row.PartialIndexUpdateHelper
-	if n := len(d.run.td.tableDesc().PartialIndexes()); n > 0 {
-		err := pm.Init(nil /* partialIndexPutVals */, sourceVals[:n], d.run.td.tableDesc())
+	if n := len(r.td.tableDesc().PartialIndexes()); n > 0 {
+		err := pm.Init(nil /* partialIndexPutVals */, sourceVals[:n], r.td.tableDesc())
 		if err != nil {
 			return err
 		}
@@ -166,53 +181,52 @@ func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums)
 	// Keep track of the vector index partitions to update. This information is
 	// passed to tableInserter.row below.
 	var vh row.VectorIndexUpdateHelper
-	if n := len(d.run.td.tableDesc().VectorIndexes()); n > 0 {
-		vh.InitForDel(sourceVals[:n], d.run.td.tableDesc())
+	if n := len(r.td.tableDesc().VectorIndexes()); n > 0 {
+		vh.InitForDel(sourceVals[:n], r.td.tableDesc())
 	}
 
 	// Queue the deletion in the KV batch.
-	if err := d.run.td.row(
-		params.ctx, deleteVals, pm, vh, false /* mustValidateOldPKValues */, d.run.traceKV,
+	if err := r.td.row(
+		params.ctx, deleteVals, pm, vh, false /* mustValidateOldPKValues */, r.traceKV,
 	); err != nil {
 		return err
 	}
 
 	// If result rows need to be accumulated, do it.
-	if d.run.td.rows != nil {
+	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.
 		//
-		// d.run.rows.NumCols() is guaranteed to only contain the requested
+		// r.rows.NumCols() is guaranteed to only contain the requested
 		// public columns.
-		resultValues := make(tree.Datums, d.run.td.rows.NumCols())
 		largestRetIdx := -1
-		for i := range d.run.rowIdxToRetIdx {
-			retIdx := d.run.rowIdxToRetIdx[i]
+		for i := range r.rowIdxToRetIdx {
+			retIdx := r.rowIdxToRetIdx[i]
 			if retIdx >= 0 {
 				if retIdx >= largestRetIdx {
 					largestRetIdx = retIdx
 				}
-				resultValues[retIdx] = deleteVals[i]
+				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 d.run.numPassthrough > 0 {
-			passthroughBegin := len(d.run.td.rd.FetchCols)
-			passthroughEnd := passthroughBegin + d.run.numPassthrough
+		if r.numPassthrough > 0 {
+			passthroughBegin := len(r.td.rd.FetchCols)
+			passthroughEnd := passthroughBegin + r.numPassthrough
 			passthroughValues := deleteVals[passthroughBegin:passthroughEnd]
 
-			for i := 0; i < d.run.numPassthrough; i++ {
+			for i := 0; i < r.numPassthrough; i++ {
 				largestRetIdx++
-				resultValues[largestRetIdx] = passthroughValues[i]
+				r.resultRowBuffer[largestRetIdx] = passthroughValues[i]
 			}
 
 		}
 
-		if _, err := d.run.td.rows.AddRow(params.ctx, resultValues); err != nil {
+		if _, err := r.td.rows.AddRow(params.ctx, r.resultRowBuffer); err != nil {
 			return err
 		}
 	}

pkg/sql/update.go

@@ -48,6 +48,10 @@ type updateRun struct {
 	// BatchedNext() has completed the work already.
 	done bool
 
+	// resultRowBuffer is used to prepare a result row for accumulation
+	// into the row container above, when rowsNeeded is set.
+	resultRowBuffer tree.Datums
+
 	// traceKV caches the current KV tracing flag.
 	traceKV bool
 
@@ -68,16 +72,26 @@ type updateRun struct {
 	regionLocalInfo regionLocalInfoType
 }
 
+func (r *updateRun) initRowContainer(params runParams, columns colinfo.ResultColumns) {
+	if !r.rowsNeeded {
+		return
+	}
+	r.tu.rows = rowcontainer.NewRowContainer(
+		params.p.Mon().MakeBoundAccount(),
+		colinfo.ColTypeInfoFromResCols(columns),
+	)
+	r.resultRowBuffer = make([]tree.Datum, len(columns))
+	for i := range r.resultRowBuffer {
+		r.resultRowBuffer[i] = tree.DNull
+	}
+}
+
 func (u *updateNode) startExec(params runParams) error {
 	// cache traceKV during execution, to avoid re-evaluating it for every row.
 	u.run.traceKV = params.p.ExtendedEvalContext().Tracing.KVTracingEnabled()
 
-	if u.run.rowsNeeded {
-		u.run.tu.rows = rowcontainer.NewRowContainer(
-			params.p.Mon().MakeBoundAccount(),
-			colinfo.ColTypeInfoFromResCols(u.columns),
-		)
-	}
+	u.run.initRowContainer(params, u.columns)
+
 	return u.run.tu.init(params.ctx, params.p.txn, params.EvalContext())
 }
 
@@ -118,7 +132,7 @@ func (u *updateNode) BatchedNext(params runParams) (bool, error) {
 
 		// Process the update for the current input row, potentially
 		// accumulating the result row for later.
-		if err := u.processSourceRow(params, u.input.Values()); err != nil {
+		if err := u.run.processSourceRow(params, u.input.Values()); err != nil {
 			return false, err
 		}
 
@@ -156,7 +170,7 @@ func (u *updateNode) BatchedNext(params runParams) (bool, error) {
 
 // processSourceRow processes one row from the source for update and, if
 // result rows are needed, saves it in the result row container.
-func (u *updateNode) processSourceRow(params runParams, sourceVals tree.Datums) error {
+func (r *updateRun) processSourceRow(params runParams, sourceVals tree.Datums) error {
 	// sourceVals contains values for the columns from the table, in the order of the
 	// table descriptor. (One per column in u.tw.ru.FetchCols)
 	//
@@ -166,42 +180,42 @@ func (u *updateNode) processSourceRow(params runParams, sourceVals tree.Datums)
 	// oldValues is the prefix of sourceVals that corresponds to real
 	// stored columns in the table, that is, excluding the RHS assignment
 	// expressions.
-	oldValues := sourceVals[:len(u.run.tu.ru.FetchCols)]
+	oldValues := sourceVals[:len(r.tu.ru.FetchCols)]
 	sourceVals = sourceVals[len(oldValues):]
 
 	// The update values follow the fetch values and their order corresponds to the order of ru.UpdateCols.
-	updateValues := sourceVals[:len(u.run.tu.ru.UpdateCols)]
+	updateValues := sourceVals[:len(r.tu.ru.UpdateCols)]
 	sourceVals = sourceVals[len(updateValues):]
 
 	// The passthrough values follow the update values.
-	passthroughValues := sourceVals[:u.run.numPassthrough]
+	passthroughValues := sourceVals[:r.numPassthrough]
 	sourceVals = sourceVals[len(passthroughValues):]
 
 	// Verify the schema constraints. For consistency with INSERT/UPSERT
 	// and compatibility with PostgreSQL, we must do this before
 	// processing the CHECK constraints.
-	if err := enforceNotNullConstraints(updateValues, u.run.tu.ru.UpdateCols); err != nil {
+	if err := enforceNotNullConstraints(updateValues, r.tu.ru.UpdateCols); err != nil {
 		return err
 	}
 
 	// Run the CHECK constraints, if any. CheckHelper will either evaluate the
 	// constraints itself, or else inspect boolean columns from the input that
 	// contain the results of evaluation.
-	if !u.run.checkOrds.Empty() {
+	if !r.checkOrds.Empty() {
 		if err := checkMutationInput(
 			params.ctx, params.EvalContext(), &params.p.semaCtx, params.p.SessionData(),
-			u.run.tu.tableDesc(), u.run.checkOrds, sourceVals[:u.run.checkOrds.Len()],
+			r.tu.tableDesc(), r.checkOrds, sourceVals[:r.checkOrds.Len()],
 		); err != nil {
 			return err
 		}
-		sourceVals = sourceVals[u.run.checkOrds.Len():]
+		sourceVals = sourceVals[r.checkOrds.Len():]
 	}
 
 	// Create a set of partial index IDs to not add entries or remove entries
 	// from. Put values are followed by del values.
 	var pm row.PartialIndexUpdateHelper
-	if n := len(u.run.tu.tableDesc().PartialIndexes()); n > 0 {
-		err := pm.Init(sourceVals[:n], sourceVals[n:n*2], u.run.tu.tableDesc())
+	if n := len(r.tu.tableDesc().PartialIndexes()); n > 0 {
+		err := pm.Init(sourceVals[:n], sourceVals[n:n*2], r.tu.tableDesc())
 		if err != nil {
 			return err
 		}
@@ -213,54 +227,53 @@ func (u *updateNode) processSourceRow(params runParams, sourceVals tree.Datums)
 	// Order of column values is put partitions, quantized vectors, followed by
 	// del partitions
 	var vh row.VectorIndexUpdateHelper
-	if n := len(u.run.tu.tableDesc().VectorIndexes()); n > 0 {
-		vh.InitForPut(sourceVals[:n], sourceVals[n:n*2], u.run.tu.tableDesc())
-		vh.InitForDel(sourceVals[n*2:n*3], u.run.tu.tableDesc())
+	if n := len(r.tu.tableDesc().VectorIndexes()); n > 0 {
+		vh.InitForPut(sourceVals[:n], sourceVals[n:n*2], r.tu.tableDesc())
+		vh.InitForDel(sourceVals[n*2:n*3], r.tu.tableDesc())
 	}
 
 	// Error out the update if the enforce_home_region session setting is on and
 	// the row's locality doesn't match the gateway region.
-	if err := u.run.regionLocalInfo.checkHomeRegion(updateValues); err != nil {
+	if err := r.regionLocalInfo.checkHomeRegion(updateValues); err != nil {
 		return err
 	}
 
 	// Queue the insert in the KV batch.
-	newValues, err := u.run.tu.rowForUpdate(
-		params.ctx, oldValues, updateValues, pm, vh, false /* mustValidateOldPKValues */, u.run.traceKV,
+	newValues, err := r.tu.rowForUpdate(
+		params.ctx, oldValues, updateValues, pm, vh, false /* mustValidateOldPKValues */, r.traceKV,
 	)
 	if err != nil {
 		return err
 	}
 
 	// If result rows need to be accumulated, do it.
-	if u.run.tu.rows != nil {
+	if r.tu.rows != nil {
 		// The new values can include all columns,  so the values may contain
 		// additional columns for every newly added column not yet visible. We do
 		// not want them to be available for RETURNING.
 		//
 		// MakeUpdater guarantees that the first columns of the new values
 		// are those specified u.columns.
-		resultValues := make([]tree.Datum, len(u.columns))
 		largestRetIdx := -1
-		for i := range u.run.rowIdxToRetIdx {
-			retIdx := u.run.rowIdxToRetIdx[i]
+		for i := range r.rowIdxToRetIdx {
+			retIdx := r.rowIdxToRetIdx[i]
 			if retIdx >= 0 {
 				if retIdx >= largestRetIdx {
 					largestRetIdx = retIdx
 				}
-				resultValues[retIdx] = newValues[i]
+				r.resultRowBuffer[retIdx] = newValues[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 FROM clause of the update).
-		for i := 0; i < u.run.numPassthrough; i++ {
+		for i := 0; i < r.numPassthrough; i++ {
 			largestRetIdx++
-			resultValues[largestRetIdx] = passthroughValues[i]
+			r.resultRowBuffer[largestRetIdx] = passthroughValues[i]
 		}
 
-		if _, err := u.run.tu.rows.AddRow(params.ctx, resultValues); err != nil {
+		if _, err := r.tu.rows.AddRow(params.ctx, r.resultRowBuffer); err != nil {
 			return err
 		}
 	}