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: #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. Requirements to hit the bug are:

1. A table with multiple column families.
2. A column family containing a single PK column.
3. That column family is not the first column family.
4. That column family existed before its column was in the PK.
5. That column must be of type FLOAT4/8, DECIMAL, JSON, collated string
   type, or array.
6. An update that changes that column from a composite value to a non-
   composite value.

pkg/sql/colenc/encode.go

@@ -701,7 +701,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

@@ -208,25 +208,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
 			}