sql/row: fix updates of single-composite-column families

When updating a single-column family which contains what could be a
composite value from the primary key, we still need to issue a Del even
if the new value for the column is not composite, because the previous
value might have been composite.

Fixes: cockroachdb#131860

Release note (bug fix): Fix a rare bug in which an update of a primary
key column which is also the only column in a separate column family can
sometimes fail to update the primary index. This bug has existed since
v22.2.

pkg/sql/colenc/encode.go

@@ -708,7 +708,7 @@ func (b *BatchEncoder) skipColumnNotInPrimaryIndexValue(
 	colID catid.ColumnID, vec coldata.Vec, row int,
 ) bool {
 	// Reuse this function but fake out the value and handle composites here.
-	if skip := b.rh.SkipColumnNotInPrimaryIndexValue(colID, tree.DNull); skip {
+	if skip, _ := b.rh.SkipColumnNotInPrimaryIndexValue(colID, tree.DNull); skip {
 		if !b.compositeColumnIDs.Contains(int(colID)) {
 			return true
 		}

pkg/sql/logictest/testdata/logic_test/column_families

@@ -79,3 +79,30 @@ ORDER BY message
 fetched: /t/t_pkey/1/2.00/x -> /1.00
 fetched: /t/t_pkey/1/2/y -> /2.00
 fetched: /t/t_pkey/1/2/z -> /1
+
+# Regression test for #131860.
+
+statement ok
+CREATE TABLE abc (a INT NOT NULL, b FLOAT NOT NULL, c INT, FAMILY (a), FAMILY (b), FAMILY (c))
+
+statement ok
+INSERT INTO abc VALUES (4, -0, 6)
+
+statement ok
+ALTER TABLE abc ADD PRIMARY KEY (a, b)
+
+statement ok
+UPDATE abc SET c = NULL WHERE a = 4 AND b = -0
+
+query IFI
+SELECT * FROM abc
+----
+4  -0  NULL
+
+statement ok
+UPDATE abc SET b = 0 WHERE a = 4 AND b = -0;
+
+query IFI
+SELECT * FROM abc
+----
+4  0  NULL

pkg/sql/row/helper.go

@@ -206,25 +206,26 @@ func (rh *RowHelper) encodeSecondaryIndexes(
 // SkipColumnNotInPrimaryIndexValue returns true if the value at column colID
 // does not need to be encoded, either because it is already part of the primary
 // key, or because it is not part of the primary index altogether. Composite
-// datums are considered too, so a composite datum in a PK will return false.
+// datums are considered too, so a composite datum in a PK will return false
+// (but will return true for couldBeComposite).
 func (rh *RowHelper) SkipColumnNotInPrimaryIndexValue(
 	colID descpb.ColumnID, value tree.Datum,
-) bool {
+) (skip, couldBeComposite bool) {
 	if rh.primaryIndexKeyCols.Empty() {
 		rh.primaryIndexKeyCols = rh.TableDesc.GetPrimaryIndex().CollectKeyColumnIDs()
 		rh.primaryIndexValueCols = rh.TableDesc.GetPrimaryIndex().CollectPrimaryStoredColumnIDs()
 	}
 	if !rh.primaryIndexKeyCols.Contains(colID) {
-		return !rh.primaryIndexValueCols.Contains(colID)
+		return !rh.primaryIndexValueCols.Contains(colID), false
 	}
 	if cdatum, ok := value.(tree.CompositeDatum); ok {
 		// Composite columns are encoded in both the key and the value.
-		return !cdatum.IsComposite()
+		return !cdatum.IsComposite(), true
 	}
 	// Skip primary key columns as their values are encoded in the key of
 	// each family. Family 0 is guaranteed to exist and acts as a
 	// sentinel.
-	return true
+	return true, false
 }
 
 func (rh *RowHelper) SortedColumnFamily(famID descpb.FamilyID) ([]descpb.ColumnID, bool) {

pkg/sql/row/writer.go

@@ -134,15 +134,25 @@ func prepareInsertOrUpdateBatch(
 			if !ok {
 				continue
 			}
+
+			var marshaled roachpb.Value
+			var err error
+			typ := fetchedCols[idx].GetType()
+
 			// Skip any values with a default ID not stored in the primary index,
 			// which can happen if we are adding new columns.
-			if skip := helper.SkipColumnNotInPrimaryIndexValue(family.DefaultColumnID, values[idx]); skip {
-				continue
-			}
-			typ := fetchedCols[idx].GetType()
-			marshaled, err := valueside.MarshalLegacy(typ, values[idx])
-			if err != nil {
-				return nil, err
+			skip, couldBeComposite := helper.SkipColumnNotInPrimaryIndexValue(family.DefaultColumnID, values[idx])
+			if skip {
+				// If the column could be composite, there could be a previous KV, so we
+				// still need to issue a Delete.
+				if !couldBeComposite {
+					continue
+				}
+			} else {
+				marshaled, err = valueside.MarshalLegacy(typ, values[idx])
+				if err != nil {
+					return nil, err
+				}
 			}
 
 			if marshaled.RawBytes == nil {
@@ -178,7 +188,7 @@ func prepareInsertOrUpdateBatch(
 				continue
 			}
 
-			if skip := helper.SkipColumnNotInPrimaryIndexValue(colID, values[idx]); skip {
+			if skip, _ := helper.SkipColumnNotInPrimaryIndexValue(colID, values[idx]); skip {
 				continue
 			}