sql: dedup primary index value encoding

pkg/sql/row/deleter.go

@@ -15,7 +15,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
-	"github.com/cockroachdb/cockroach/pkg/sql/rowenc/valueside"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
@@ -254,47 +253,16 @@ func (rd *Deleter) DeleteRow(
 	return nil
 }
 
-// encodeValueForPrimaryIndexFamily encodes the expected roachpb.Value
-// for the given family and valuses.
-//
-// TODO(ssd): Lots of duplication between this and
-// prepareInsertOrUpdateBatch. This is rather unfortunate.
+// encodeValueForPrimaryIndexFamily encodes the expected roachpb.Value for
+// the given family and values. It is used to compute the expected previous
+// value for CPut-based deletes (origin-timestamp logical replication and
+// mustValidateOldPKValues paths).
 func (rd *Deleter) encodeValueForPrimaryIndexFamily(
 	family *descpb.ColumnFamilyDescriptor, values []tree.Datum,
 ) (roachpb.Value, error) {
-	if len(family.ColumnIDs) == 1 && family.ColumnIDs[0] == family.DefaultColumnID && family.ID != 0 {
-		idx, ok := rd.FetchColIDtoRowIndex.Get(family.DefaultColumnID)
-		if !ok {
-			return roachpb.Value{}, nil
-		}
-		if skip, _ := rd.Helper.SkipColumnNotInPrimaryIndexValue(family.DefaultColumnID, values[idx]); skip {
-			return roachpb.Value{}, nil
-		}
-		typ := rd.FetchCols[idx].GetType()
-		marshaled, err := valueside.MarshalLegacy(typ, values[idx])
-		if err != nil {
-			return roachpb.Value{}, err
-		}
-
-		return marshaled, err
-	}
-
-	rd.rawValueBuf = rd.rawValueBuf[:0]
-	familySortedColumnIDs, ok := rd.Helper.SortedColumnFamily(family.ID)
-	if !ok {
-		return roachpb.Value{}, errors.AssertionFailedf("invalid family sorted column id map")
-	}
-
-	var err error
-	rd.rawValueBuf, err = rd.Helper.encodePrimaryIndexValuesToBuf(values, rd.FetchColIDtoRowIndex, familySortedColumnIDs, rd.FetchCols, rd.rawValueBuf)
-	if err != nil {
-		return roachpb.Value{}, err
-	}
-	ret := roachpb.Value{}
-	// For family 0, we expect a value even when no columns have
-	// been encoded to oldBytes.
-	if family.ID == 0 || len(rd.rawValueBuf) > 0 {
-		ret.SetTuple(rd.rawValueBuf)
-	}
-	return ret, nil
+	res, retBuf, err := rd.Helper.PrimaryIndexEncoder().EncodeFamily(
+		family, rd.FetchColIDtoRowIndex, values, rd.rawValueBuf,
+	)
+	rd.rawValueBuf = retBuf
+	return res.Value, err
 }

pkg/sql/row/helper.go

@@ -25,7 +25,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc/rowencpb"
-	"github.com/cockroachdb/cockroach/pkg/sql/rowenc/valueside"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/idxtype"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -125,11 +124,13 @@ type RowHelper struct {
 		secondary [][]encoding.Direction
 	}
 
-	// Computed and cached.
+	// PrimaryIndexKeyPrefix is the encoded key prefix for the table's
+	// primary index. Initialized by Init.
 	PrimaryIndexKeyPrefix []byte
-	primaryIndexKeyCols   catalog.TableColSet
-	primaryIndexValueCols catalog.TableColSet
-	sortedColumnFamilies  map[descpb.FamilyID][]descpb.ColumnID
+	// pkEncoder owns the per-table state for encoding primary-index
+	// family values (key/stored column sets, sorted family column IDs).
+	// Initialized by Init and reused across every row.
+	pkEncoder rowenc.PrimaryIndexEncoder
 
 	// Used to build tmpTombstones for non-Serializable uniqueness checks.
 	index2UniqueWithTombstoneEntry map[catalog.Index]*uniqueWithTombstoneEntry
@@ -235,6 +236,13 @@ func (rh *RowHelper) Init() {
 	rh.PrimaryIndexKeyPrefix = rowenc.MakeIndexKeyPrefix(
 		rh.Codec, rh.TableDesc.GetID(), rh.TableDesc.GetPrimaryIndexID(),
 	)
+	rh.pkEncoder = rowenc.MakePrimaryIndexEncoder(rh.TableDesc, rh.TableDesc.GetPrimaryIndex())
+}
+
+// PrimaryIndexEncoder returns the encoder for the table's primary index.
+// Init must have been called.
+func (rh *RowHelper) PrimaryIndexEncoder() *rowenc.PrimaryIndexEncoder {
+	return &rh.pkEncoder
 }
 
 // encodePrimaryIndexKey encodes the primary index key.
@@ -403,70 +411,25 @@ func (rh *RowHelper) encodeSecondaryIndexes(
 	return rh.indexEntries, nil
 }
 
-// encodePrimaryIndexValuesToBuf encodes the given values, writing
-// into the given buffer.
-func (rh *RowHelper) encodePrimaryIndexValuesToBuf(
-	vals []tree.Datum,
-	valColIDMapping catalog.TableColMap,
-	sortedColumnIDs []descpb.ColumnID,
-	fetchedCols []catalog.Column,
-	buf []byte,
-) ([]byte, error) {
-	var lastColID descpb.ColumnID
-	for _, colID := range sortedColumnIDs {
-		idx, ok := valColIDMapping.Get(colID)
-		if !ok || vals[idx] == tree.DNull {
-			// Column not being updated or inserted.
-			continue
-		}
-
-		if skip, _ := rh.SkipColumnNotInPrimaryIndexValue(colID, vals[idx]); skip {
-			continue
-		}
-
-		col := fetchedCols[idx]
-		if lastColID > col.GetID() {
-			return nil, errors.AssertionFailedf("cannot write column id %d after %d", col.GetID(), lastColID)
-		}
-		colIDDelta := valueside.MakeColumnIDDelta(lastColID, col.GetID())
-		lastColID = col.GetID()
-		var err error
-		buf, err = valueside.Encode(buf, colIDDelta, vals[idx])
-		if err != nil {
-			return nil, err
-		}
-	}
-	return buf, nil
-}
-
 // 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
-// (but will return true for couldBeComposite).
+// 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 (but will return true for couldBeComposite).
 func (rh *RowHelper) SkipColumnNotInPrimaryIndexValue(
 	colID descpb.ColumnID, value tree.Datum,
 ) (skip, couldBeComposite bool) {
-	if rh.primaryIndexKeyCols.Empty() {
-		rh.primaryIndexKeyCols = rh.TableDesc.GetPrimaryIndex().CollectKeyColumnIDs()
-		rh.primaryIndexValueCols = rh.TableDesc.GetPrimaryIndex().CollectPrimaryStoredColumnIDs()
-	}
-	return rowenc.SkipColumnNotInPrimaryIndexValue(colID, value, rh.primaryIndexKeyCols, rh.primaryIndexValueCols)
+	// We pass an empty colMap; the encoder only consults it to compute the
+	// stored column set when the primary index uses a pre-storedcols
+	// version. RowHelper has no colMap of its own and historically did not
+	// handle that case either.
+	return rh.pkEncoder.SkipColumn(colID, value, catalog.TableColMap{})
 }
 
+// SortedColumnFamily returns the cached sorted column IDs for family famID
+// of the primary index, or (nil, false) if no such family exists.
 func (rh *RowHelper) SortedColumnFamily(famID descpb.FamilyID) ([]descpb.ColumnID, bool) {
-	if rh.sortedColumnFamilies == nil {
-		rh.sortedColumnFamilies = make(map[descpb.FamilyID][]descpb.ColumnID, rh.TableDesc.NumFamilies())
-
-		_ = rh.TableDesc.ForeachFamily(func(family *descpb.ColumnFamilyDescriptor) error {
-			colIDs := append([]descpb.ColumnID{}, family.ColumnIDs...)
-			sort.Sort(descpb.ColumnIDs(colIDs))
-			rh.sortedColumnFamilies[family.ID] = colIDs
-			return nil
-		})
-	}
-	colIDs, ok := rh.sortedColumnFamilies[famID]
-	return colIDs, ok
+	return rh.pkEncoder.SortedFamilyColumnIDs(famID)
 }
 
 // CheckRowSize checks the row size against the max_row_size limits. Over

pkg/sql/row/inserter.go

@@ -199,7 +199,7 @@ func (ri *Inserter) InsertRow(
 
 	// Add the new values to the primary index.
 	ri.valueBuf, err = prepareInsertOrUpdateBatch(
-		ctx, b, &ri.Helper, primaryIndexKey, ri.InsertCols, values, ri.InsertColIDtoRowIndex,
+		ctx, b, &ri.Helper, primaryIndexKey, values, ri.InsertColIDtoRowIndex,
 		ri.InsertColIDtoRowIndex, &ri.key, &ri.value, ri.valueBuf, oth, nil, /* oldValues */
 		kvOp, false /* mustValidateOldPKValues */, traceKV,
 	)

pkg/sql/row/updater.go

@@ -420,7 +420,7 @@ func (ru *Updater) UpdateRow(
 		kvOp = PutOp
 	}
 	ru.valueBuf, err = prepareInsertOrUpdateBatch(
-		ctx, b, &ru.Helper, primaryIndexKey, ru.FetchCols, ru.newValues, ru.FetchColIDtoRowIndex,
+		ctx, b, &ru.Helper, primaryIndexKey, ru.newValues, ru.FetchColIDtoRowIndex,
 		ru.UpdateColIDtoRowIndex, &ru.key, &ru.value, ru.valueBuf, oth, oldValues,
 		kvOp, mustValidateOldPKValues, traceKV,
 	)