Merge #143778

143778: row, sql: add CPut of primary index to Updater and Deleter r=yuzefovich a=michae2

The upcoming fast paths for UPDATE and DELETE will skip the initial scan
of the target row. Instead, they will rely on CPut(s) to the primary
index to confirm that the target row exists and matches the WHERE clause
of the statement.

This PR adds a new CPut path to the Updater and Deleter which uses
expValue to validate the old values. It closely mirrors the CPut paths
added for the OriginTimestampCPutHelper in #130512. The main difference
is that unlike the OriginTimestampCPutHelper, these paths do not need to
delete all-NULL non-0 column families when not overwriting.

Only writes to the primary index need validation using CPut with
expValue. Writes to secondary indexes should be unchanged. If the CPut
to the primary index fails, all writes for the statement will be rolled
back using a savepoint (next PR).

Informs: #71153

Release note: None

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

Author: craig[bot]

pkg/crosscluster/logical/lww_kv_processor.go

@@ -503,7 +503,10 @@ func (p *kvTableWriter) updateRow(
 		// and destination clusters.
 		// ShouldWinTie: true,
 	}
-	_, err := p.ru.UpdateRow(ctx, b, p.oldVals, p.newVals, ph, vh, oth, false)
+	_, err := p.ru.UpdateRow(
+		ctx, b, p.oldVals, p.newVals, ph, vh, oth, false, /* mustValidateOldPKValues */
+		false, /* traceKV */
+	)
 	return err
 }
 
@@ -526,7 +529,9 @@ func (p *kvTableWriter) deleteRow(
 		// ShouldWinTie: true,
 	}
 
-	return p.rd.DeleteRow(ctx, b, p.oldVals, ph, vh, oth, false)
+	return p.rd.DeleteRow(
+		ctx, b, p.oldVals, ph, vh, oth, false /* mustValidateOldPKValues */, false, /* traceKV */
+	)
 }
 
 func (p *kvTableWriter) fillOld(vals cdcevent.Row) error {

pkg/crosscluster/logical/tombstone_updater.go

@@ -155,6 +155,7 @@ func (tu *tombstoneUpdater) addToBatch(
 			OriginTimestamp:    afterRow.MvccTimestamp,
 			PreviousWasDeleted: true,
 		},
+		false, /* mustValidateOldPKValues */
 		false, /* traceKV */
 	)
 }

pkg/sql/backfill/backfill.go

@@ -402,7 +402,7 @@ func (cb *ColumnBackfiller) RunColumnBackfillChunk(
 		var pm row.PartialIndexUpdateHelper
 		var vh row.VectorIndexUpdateHelper
 		if _, err := ru.UpdateRow(
-			ctx, b, oldValues, updateValues, pm, vh, nil, traceKV,
+			ctx, b, oldValues, updateValues, pm, vh, nil, false /* mustValidateOldPKValues */, traceKV,
 		); err != nil {
 			return roachpb.Key{}, err
 		}

pkg/sql/delete.go

@@ -171,7 +171,9 @@ func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums)
 	}
 
 	// Queue the deletion in the KV batch.
-	if err := d.run.td.row(params.ctx, deleteVals, pm, vh, d.run.traceKV); err != nil {
+	if err := d.run.td.row(
+		params.ctx, deleteVals, pm, vh, false /* mustValidateOldPKValues */, d.run.traceKV,
+	); err != nil {
 		return err
 	}
 

pkg/sql/row/deleter.go

@@ -141,6 +141,7 @@ func (rd *Deleter) DeleteRow(
 	pm PartialIndexUpdateHelper,
 	vh VectorIndexUpdateHelper,
 	oth *OriginTimestampCPutHelper,
+	mustValidateOldPKValues bool,
 	traceKV bool,
 ) error {
 	b := &KVBatchAdapter{Batch: batch}
@@ -197,7 +198,7 @@ func (rd *Deleter) DeleteRow(
 		familyID := family.ID
 		rd.key = keys.MakeFamilyKey(primaryIndexKey, uint32(familyID))
 
-		if oth.IsSet() {
+		if oth.IsSet() || mustValidateOldPKValues {
 			var expValue []byte
 			if !oth.PreviousWasDeleted {
 				prevValue, err := rd.encodeValueForPrimaryIndexFamily(family, values)
@@ -208,7 +209,11 @@ func (rd *Deleter) DeleteRow(
 					expValue = prevValue.TagAndDataBytes()
 				}
 			}
-			oth.DelWithCPut(ctx, b, &rd.key, expValue, traceKV)
+			if oth.IsSet() {
+				oth.DelWithCPut(ctx, b, &rd.key, expValue, traceKV)
+			} else {
+				delWithCPutFn(ctx, b, &rd.key, expValue, traceKV, rd.Helper.primIndexValDirs)
+			}
 		} else {
 			delFn(ctx, b, &rd.key, !rd.primaryLocked /* needsLock */, traceKV, rd.Helper.primIndexValDirs)
 		}

pkg/sql/row/helper.go

@@ -494,6 +494,24 @@ func delFn(
 	}
 }
 
+func delWithCPutFn(
+	ctx context.Context,
+	b Putter,
+	key *roachpb.Key,
+	expVal []byte,
+	traceKV bool,
+	keyEncodingDirs []encoding.Direction,
+) {
+	if traceKV {
+		if keyEncodingDirs != nil {
+			log.VEventf(ctx, 2, "CPut %s -> nil (delete)", keys.PrettyPrint(keyEncodingDirs, *key))
+		} else {
+			log.VEventf(ctx, 2, "CPut %s -> nil (delete)", *key)
+		}
+	}
+	b.CPut(key, nil, expVal)
+}
+
 func (rh *RowHelper) deleteIndexEntry(
 	ctx context.Context,
 	b Putter,
@@ -550,7 +568,10 @@ func (oh *OriginTimestampCPutHelper) CPutFn(
 	traceKV bool,
 ) {
 	if traceKV {
-		log.VEventfDepth(ctx, 1, 2, "CPutWithOriginTimestamp %s -> %s @ %s", *key, value.PrettyPrint(), oh.OriginTimestamp)
+		log.VEventfDepth(
+			ctx, 1, 2, "CPutWithOriginTimestamp %s -> %s (swap) @ %s", *key, value.PrettyPrint(),
+			oh.OriginTimestamp,
+		)
 	}
 	b.CPutWithOriginTimestamp(key, value, expVal, oh.OriginTimestamp)
 }
@@ -559,7 +580,9 @@ func (oh *OriginTimestampCPutHelper) DelWithCPut(
 	ctx context.Context, b Putter, key *roachpb.Key, expVal []byte, traceKV bool,
 ) {
 	if traceKV {
-		log.VEventfDepth(ctx, 1, 2, "CPutWithOriginTimestamp %s -> nil (delete) @ %s", key, oh.OriginTimestamp)
+		log.VEventfDepth(
+			ctx, 1, 2, "CPutWithOriginTimestamp %s -> nil (delete) @ %s", key, oh.OriginTimestamp,
+		)
 	}
 	b.CPutWithOriginTimestamp(key, nil, expVal, oh.OriginTimestamp)
 }

pkg/sql/row/inserter.go

@@ -199,7 +199,7 @@ func (ri *Inserter) InsertRow(
 	ri.valueBuf, err = prepareInsertOrUpdateBatch(
 		ctx, b, &ri.Helper, primaryIndexKey, ri.InsertCols, values, ri.InsertColIDtoRowIndex,
 		ri.InsertColIDtoRowIndex, &ri.key, &ri.value, ri.valueBuf, oth, nil, /* oldValues */
-		kvOp, traceKV,
+		kvOp, false /* mustValidateOldPKValues */, traceKV,
 	)
 	if err != nil {
 		return err

pkg/sql/row/updater.go

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/deduplicate"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/intsets"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 )
 
@@ -241,6 +242,7 @@ func (ru *Updater) UpdateRow(
 	pm PartialIndexUpdateHelper,
 	vh VectorIndexUpdateHelper,
 	oth *OriginTimestampCPutHelper,
+	mustValidateOldPKValues bool,
 	traceKV bool,
 ) ([]tree.Datum, error) {
 	if len(oldValues) != len(ru.FetchCols) {
@@ -396,7 +398,9 @@ func (ru *Updater) UpdateRow(
 		// followed by a CPut we could skip the Del altogether. Furthermore, if
 		// we acquired the lock on this index during the initial scan we could
 		// replace a CPut with a Put.
-		if err := ru.rd.DeleteRow(ctx, batch, oldValues, pm, vh, oth, traceKV); err != nil {
+		if err := ru.rd.DeleteRow(
+			ctx, batch, oldValues, pm, vh, oth, mustValidateOldPKValues, traceKV,
+		); err != nil {
 			return nil, err
 		}
 		if err := ru.ri.InsertRow(
@@ -418,7 +422,7 @@ func (ru *Updater) UpdateRow(
 	ru.valueBuf, err = prepareInsertOrUpdateBatch(
 		ctx, b, &ru.Helper, primaryIndexKey, ru.FetchCols, ru.newValues, ru.FetchColIDtoRowIndex,
 		ru.UpdateColIDtoRowIndex, &ru.key, &ru.value, ru.valueBuf, oth, oldValues,
-		kvOp, traceKV,
+		kvOp, mustValidateOldPKValues, traceKV,
 	)
 	if err != nil {
 		return nil, err
@@ -658,3 +662,21 @@ func (ru *Updater) IsColumnOnlyUpdate() bool {
 	// operations) and these should be split.
 	return !ru.primaryKeyColChange && ru.DeleteHelper == nil && len(ru.Helper.Indexes) == 0
 }
+
+func updateCPutFn(
+	ctx context.Context,
+	b Putter,
+	key *roachpb.Key,
+	value *roachpb.Value,
+	expVal []byte,
+	traceKV bool,
+	keyEncodingDirs []encoding.Direction,
+) {
+	if traceKV {
+		log.VEventfDepth(
+			ctx, 1, 2, "CPut %s -> %s (swap)", keys.PrettyPrint(keyEncodingDirs, *key),
+			value.PrettyPrint(),
+		)
+	}
+	b.CPut(key, value, expVal)
+}

pkg/sql/row/writer.go

@@ -79,6 +79,8 @@ func ColMapping(fromCols, toCols []catalog.Column) []int {
 //     capacity to avoid allocations. The function returns the slice.
 //   - kvOp indicates which KV write operation should be used. If it is PutOp,
 //     it also indicates that the old keys have been locked.
+//   - mustValidateOldPKValues indicates whether the expected previous row must
+//     be verified (using CPut)
 //   - traceKV is to be set to log the KV operations added to the batch.
 func prepareInsertOrUpdateBatch(
 	ctx context.Context,
@@ -95,6 +97,7 @@ func prepareInsertOrUpdateBatch(
 	oth *OriginTimestampCPutHelper,
 	oldValues []tree.Datum,
 	kvOp KVInsertOp,
+	mustValidateOldPKValues bool,
 	traceKV bool,
 ) ([]byte, error) {
 	families := helper.TableDesc.GetFamilies()
@@ -182,7 +185,7 @@ func prepareInsertOrUpdateBatch(
 			}
 
 			var oldVal []byte
-			if oth.IsSet() && len(oldValues) > 0 {
+			if (oth.IsSet() || mustValidateOldPKValues) && len(oldValues) > 0 {
 				// If the column could be composite, we only encode the old value if it
 				// was a composite value.
 				if !couldBeComposite || oldValues[idx].(tree.CompositeDatum).IsComposite() {
@@ -204,8 +207,12 @@ func prepareInsertOrUpdateBatch(
 				} else if overwrite {
 					// If the new family contains a NULL value, then we must
 					// delete any pre-existing row.
-					needsLock := !oldKeysLocked
-					delFn(ctx, batch, kvKey, needsLock, traceKV, helper.primIndexValDirs)
+					if mustValidateOldPKValues {
+						delWithCPutFn(ctx, batch, kvKey, oldVal, traceKV, helper.primIndexValDirs)
+					} else {
+						needsLock := !oldKeysLocked
+						delFn(ctx, batch, kvKey, needsLock, traceKV, helper.primIndexValDirs)
+					}
 				}
 			} else {
 				// We only output non-NULL values. Non-existent column keys are
@@ -217,6 +224,8 @@ func prepareInsertOrUpdateBatch(
 
 				if oth.IsSet() {
 					oth.CPutFn(ctx, batch, kvKey, &marshaled, oldVal, traceKV)
+				} else if mustValidateOldPKValues {
+					updateCPutFn(ctx, batch, kvKey, &marshaled, oldVal, traceKV, helper.primIndexValDirs)
 				} else {
 					// TODO(yuzefovich): in case of multiple column families,
 					// whenever we locked the primary index during the initial
@@ -254,7 +263,7 @@ func prepareInsertOrUpdateBatch(
 		// If we are using OriginTimestamp ConditionalPuts, calculate the expected
 		// value.
 		var expBytes []byte
-		if oth.IsSet() && len(oldValues) > 0 {
+		if (oth.IsSet() || mustValidateOldPKValues) && len(oldValues) > 0 {
 			var oldBytes []byte
 			oldBytes, err = helper.encodePrimaryIndexValuesToBuf(oldValues, valColIDMapping, familySortedColumnIDs, fetchedCols, oldBytes)
 			if err != nil {
@@ -277,8 +286,12 @@ func prepareInsertOrUpdateBatch(
 			} else if overwrite {
 				// The family might have already existed but every column in it is being
 				// set to NULL, so delete it.
-				needsLock := !oldKeysLocked
-				delFn(ctx, batch, kvKey, needsLock, traceKV, helper.primIndexValDirs)
+				if mustValidateOldPKValues {
+					delWithCPutFn(ctx, batch, kvKey, expBytes, traceKV, helper.primIndexValDirs)
+				} else {
+					needsLock := !oldKeysLocked
+					delFn(ctx, batch, kvKey, needsLock, traceKV, helper.primIndexValDirs)
+				}
 			}
 		} else {
 			// Copy the contents of rawValueBuf into the roachpb.Value. This is
@@ -290,6 +303,8 @@ func prepareInsertOrUpdateBatch(
 			}
 			if oth.IsSet() {
 				oth.CPutFn(ctx, batch, kvKey, kvValue, expBytes, traceKV)
+			} else if mustValidateOldPKValues {
+				updateCPutFn(ctx, batch, kvKey, kvValue, expBytes, traceKV, helper.primIndexValDirs)
 			} else {
 				putFn(ctx, batch, kvKey, kvValue, traceKV, helper.primIndexValDirs)
 			}

pkg/sql/tablewriter_delete.go

@@ -43,6 +43,9 @@ func (td *tableDeleter) init(_ context.Context, txn *kv.Txn, evalCtx *eval.Conte
 // of the table, and are materialized only for the purpose of updating vector
 // indexes.
 //
+// The mustValidateOldPKValues parameter indicates whether the expected previous
+// row must be verified (using CPut).
+//
 // The traceKV parameter determines whether the individual K/V operations
 // should be logged to the context. We use a separate argument here instead
 // of a Value field on the context because Value access in context.Context
@@ -52,10 +55,11 @@ func (td *tableDeleter) row(
 	values tree.Datums,
 	pm row.PartialIndexUpdateHelper,
 	vh row.VectorIndexUpdateHelper,
+	mustValidateOldPKValues bool,
 	traceKV bool,
 ) error {
 	td.currentBatchSize++
-	return td.rd.DeleteRow(ctx, td.b, values, pm, vh, nil, traceKV)
+	return td.rd.DeleteRow(ctx, td.b, values, pm, vh, nil, mustValidateOldPKValues, traceKV)
 }
 
 // deleteIndex runs the kv operations necessary to delete all kv entries in the