diff --git a/pkg/sql/row/deleter.go b/pkg/sql/row/deleter.go index 66b3cfdab2b8..d01edc4a9c61 100644 --- a/pkg/sql/row/deleter.go +++ b/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 } diff --git a/pkg/sql/row/helper.go b/pkg/sql/row/helper.go index 386e447bfe28..facdefe992ec 100644 --- a/pkg/sql/row/helper.go +++ b/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 diff --git a/pkg/sql/row/inserter.go b/pkg/sql/row/inserter.go index 65bc2223ed78..dfdb73afdf2b 100644 --- a/pkg/sql/row/inserter.go +++ b/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, ) diff --git a/pkg/sql/row/updater.go b/pkg/sql/row/updater.go index 7a787faf9e6d..fea7b164a1c7 100644 --- a/pkg/sql/row/updater.go +++ b/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, ) diff --git a/pkg/sql/row/writer.go b/pkg/sql/row/writer.go index 2bc16638eb05..a22ddce06dc2 100644 --- a/pkg/sql/row/writer.go +++ b/pkg/sql/row/writer.go @@ -11,7 +11,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/keys" "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog" - "github.com/cockroachdb/cockroach/pkg/sql/rowenc/valueside" + "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/util" "github.com/cockroachdb/errors" @@ -59,34 +59,36 @@ func ColMapping(fromCols, toCols []catalog.Column) []int { return result } -// prepareInsertOrUpdateBatch constructs a KV batch that inserts or -// updates a row in KV in the primary index. +// prepareInsertOrUpdateBatch constructs a KV batch that inserts or updates +// a row in KV in the primary index. Per-family value encoding is delegated +// to rowenc.PrimaryIndexEncoder via helper.PrimaryIndexEncoder(); this +// function adds the family-key construction, KV-op dispatch (CPut / Put / +// PutMustAcquireExclusiveLock / Delete), origin-timestamp / old-PK +// validation handling, and row-size guard rails. +// // - batch is the KV batch where commands should be appended. -// - helper is the rowHelper that knows about the table being modified. +// - helper is the RowHelper that knows about the table being modified. // - primaryIndexKey is the PK prefix for the current row. -// - fetchedCols is the list of schema columns that have been fetched -// in preparation for this update. -// - values is the SQL-level row values that are being written. -// - valColIDMapping is the mapping from column IDs into positions of the slice -// values. -// - updatedColIDMapping is the mapping from column IDs into the positions of -// the updated values. -// - kvKey and kvValues must be heap-allocated scratch buffers to write -// roachpb.Key and roachpb.Value values. -// - rawValueBuf must be a scratch byte array. This must be reinitialized -// to an empty slice on each call but can be preserved at its current -// 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. +// - values is the SQL-level row values being written. +// - valColIDMapping maps column IDs to positions of the slice values. +// - updatedColIDMapping maps column IDs to positions of the updated +// values; only families that touch one of these columns are processed. +// - kvKey and kvValue are heap-allocated scratch buffers used to stage +// each family's roachpb.Key and roachpb.Value. +// - rawValueBuf is a scratch byte buffer for tuple-encoded family values; +// it can be reused across calls to amortize allocations and is returned +// for that purpose. +// - kvOp selects the KV write op for non-deletes (CPutOp / PutOp / +// PutMustAcquireExclusiveLockOp). PutOp also indicates that the old +// keys have been locked. +// - oldValues, oth, and mustValidateOldPKValues drive the construction of +// expected previous bytes for CPut-based writes. +// - traceKV enables KV-op logging. func prepareInsertOrUpdateBatch( ctx context.Context, batch Putter, helper *RowHelper, primaryIndexKey []byte, - fetchedCols []catalog.Column, values []tree.Datum, valColIDMapping catalog.TableColMap, updatedColIDMapping catalog.TableColMap, @@ -100,11 +102,11 @@ func prepareInsertOrUpdateBatch( traceKV bool, ) ([]byte, error) { families := helper.TableDesc.GetFamilies() - // TODO(ssd): We don't currently support multiple column - // families on the LDR write path. As a result, we don't have - // good end-to-end testing of multi-column family writes with - // the origin timestamp helper set. Until we write such tests, - // we error if we ever see such writes. + // TODO(ssd): We don't currently support multiple column families on + // the LDR write path. As a result, we don't have good end-to-end + // testing of multi-column family writes with the origin timestamp + // helper set. Until we write such tests, we error if we ever see such + // writes. if oth.IsSet() && len(families) > 1 { return nil, errors.AssertionFailedf("OriginTimestampCPutHelper is not yet testing with multi-column family writes") } @@ -125,197 +127,110 @@ func prepareInsertOrUpdateBatch( overwrite = true } + encoder := helper.PrimaryIndexEncoder() + wantOldValue := (oth.IsSet() || mustValidateOldPKValues) && len(oldValues) > 0 + var oldBuf []byte for i := range families { family := &families[i] - update := false - for _, colID := range family.ColumnIDs { - if _, ok := updatedColIDMapping.Get(colID); ok { - update = true - break - } - } - // We can have an empty family.ColumnIDs in the following case: - // * A table is created with the primary key not in family 0, and another column in family 0. - // * The column in family 0 is dropped, leaving the 0'th family empty. - // In this case, we must keep the empty 0'th column family in order to ensure that column family 0 - // is always encoded as the sentinel k/v for a row. - if !update && len(family.ColumnIDs) != 0 { + // Empty family.ColumnIDs occurs when the PK lives in a non-zero + // family and family 0's only column has been dropped; family 0 + // must still be visited so it emits the row sentinel. + if !familyTouchedByUpdate(family, updatedColIDMapping) && len(family.ColumnIDs) != 0 { continue } if i > 0 { - // HACK: MakeFamilyKey appends to its argument, so on every loop iteration - // after the first, trim primaryIndexKey so nothing gets overwritten. + // HACK: MakeFamilyKey appends to its argument, so on every + // loop iteration after the first, trim primaryIndexKey so + // nothing gets overwritten. // TODO(dan): Instead of this, use something like engine.ChunkAllocator. primaryIndexKey = primaryIndexKey[:len(primaryIndexKey):len(primaryIndexKey)] } - *kvKey = keys.MakeFamilyKey(primaryIndexKey, uint32(family.ID)) - // We need to ensure that column family 0 contains extra metadata, like composite primary key values. - // Additionally, the decoders expect that column family 0 is encoded with a TUPLE value tag, so we - // don't want to use the untagged value encoding. - if len(family.ColumnIDs) == 1 && family.ColumnIDs[0] == family.DefaultColumnID && family.ID != 0 { - // Storage optimization to store DefaultColumnID directly as a value. Also - // backwards compatible with the original BaseFormatVersion. - - idx, ok := valColIDMapping.Get(family.DefaultColumnID) - 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. - 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 - } - } - - var oldVal []byte - 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() { - old, err := valueside.MarshalLegacy(typ, oldValues[idx]) - if err != nil { - return nil, err - } - if old.IsPresent() { - oldVal = old.TagAndDataBytes() - } - } - } - - if !marshaled.IsPresent() { - if oth.IsSet() { - // If using OriginTimestamp'd CPuts, we _always_ want to issue a Delete - // so that we can confirm our expected bytes were correct. - oth.DelWithCPut(ctx, batch, kvKey, oldVal, traceKV) - } else if overwrite { - // If the new family contains a NULL value, then we must - // delete any pre-existing row. - if mustValidateOldPKValues { - delWithCPutFn(ctx, batch, kvKey, oldVal, traceKV, helper, primaryIndexDirs) - } else { - needsLock := !oldKeysLocked - delFn(ctx, batch, kvKey, needsLock, traceKV, helper, primaryIndexDirs) - } - } - } else { - // We only output non-NULL values. Non-existent column keys are - // considered NULL during scanning and the row sentinel ensures we know - // the row exists. - if err := helper.CheckRowSize(ctx, kvKey, marshaled.RawBytes, family.ID); err != nil { - return nil, err - } - - if oth.IsSet() { - oth.CPutFn(ctx, batch, kvKey, &marshaled, oldVal, traceKV) - } else if mustValidateOldPKValues { - updateCPutFn(ctx, batch, kvKey, &marshaled, oldVal, traceKV, helper, primaryIndexDirs) - } else { - // TODO(yuzefovich): in case of multiple column families, - // whenever we locked the primary index during the initial - // scan, we might not have locked the key for a column - // family where all columns had NULL values (because the KV - // didn't exist) and now at least one becomes non-NULL. In - // this scenario we're inserting a new KV with non-locking - // Put, yet we don't have the lock. - // - // However, at the moment we disable the lock eliding - // optimization with multiple column families, so we'll use - // the locking Put because of that. - putFn(ctx, batch, kvKey, &marshaled, traceKV, helper, primaryIndexDirs) - } - } - - continue - } - - familySortedColumnIDs, ok := helper.SortedColumnFamily(family.ID) - if !ok { - return nil, errors.AssertionFailedf("invalid family sorted column id map") - } - rawValueBuf = rawValueBuf[:0] - var err error - rawValueBuf, err = helper.encodePrimaryIndexValuesToBuf(values, valColIDMapping, familySortedColumnIDs, fetchedCols, rawValueBuf) + res, retBuf, err := encoder.EncodeFamily(family, valColIDMapping, values, rawValueBuf) if err != nil { return nil, err } + rawValueBuf = retBuf + if res.Skipped { + *kvKey = nil + continue + } - // TODO(ssd): Here and below investigate reducing the number of - // allocations required to marshal the old value. - // - // If we are using OriginTimestamp ConditionalPuts, calculate the expected - // value. var expBytes []byte - if (oth.IsSet() || mustValidateOldPKValues) && len(oldValues) > 0 { - var oldBytes []byte - oldBytes, err = helper.encodePrimaryIndexValuesToBuf(oldValues, valColIDMapping, familySortedColumnIDs, fetchedCols, oldBytes) + if wantOldValue { + // TODO(ssd): investigate reducing the number of allocations + // required to marshal the old value. + oldRes, retOldBuf, err := encoder.EncodeFamily(family, valColIDMapping, oldValues, oldBuf) if err != nil { return nil, err } - // For family 0, we expect a value even when - // no columns have been encoded to oldBytes. - if family.ID == 0 || len(oldBytes) > 0 { - old := &roachpb.Value{} - old.SetTuple(oldBytes) - expBytes = old.TagAndDataBytes() + oldBuf = retOldBuf + if oldRes.Value.IsPresent() { + expBytes = oldRes.Value.TagAndDataBytes() } } - if family.ID != 0 && len(rawValueBuf) == 0 { - if oth.IsSet() { - // If using OriginTimestamp'd CPuts, we _always_ want to issue a Delete - // so that we can confirm our expected bytes were correct. + if !res.Value.IsPresent() { + // No current value for this family. EncodeFamily guarantees + // res.Value is present for family 0, so reaching here means + // family.ID != 0. Issue a Delete in OT and overwrite paths; + // inserts skip silently. + switch { + case oth.IsSet(): + // OT'd CPuts always emit a Delete to confirm expected bytes. oth.DelWithCPut(ctx, batch, kvKey, expBytes, traceKV) - } else if overwrite { - // The family might have already existed but every column in it is being - // set to NULL, so delete it. - if mustValidateOldPKValues { - delWithCPutFn(ctx, batch, kvKey, expBytes, traceKV, helper, primaryIndexDirs) - } else { - needsLock := !oldKeysLocked - delFn(ctx, batch, kvKey, needsLock, traceKV, helper, primaryIndexDirs) - } - } - } else { - // Copy the contents of rawValueBuf into the roachpb.Value. This is - // a deep copy so rawValueBuf can be re-used by other calls to the - // function. - kvValue.SetTuple(rawValueBuf) - if err := helper.CheckRowSize(ctx, kvKey, kvValue.RawBytes, family.ID); err != nil { - return nil, err - } - if oth.IsSet() { - oth.CPutFn(ctx, batch, kvKey, kvValue, expBytes, traceKV) - } else if mustValidateOldPKValues { - updateCPutFn(ctx, batch, kvKey, kvValue, expBytes, traceKV, helper, primaryIndexDirs) - } else { - putFn(ctx, batch, kvKey, kvValue, traceKV, helper, primaryIndexDirs) + case overwrite && mustValidateOldPKValues: + delWithCPutFn(ctx, batch, kvKey, expBytes, traceKV, helper, primaryIndexDirs) + case overwrite: + delFn(ctx, batch, kvKey, !oldKeysLocked, traceKV, helper, primaryIndexDirs) } + *kvKey = nil + continue } - // Release reference to roachpb.Key. + if err := helper.CheckRowSize(ctx, kvKey, res.Value.RawBytes, family.ID); err != nil { + return nil, err + } + *kvValue = res.Value + switch { + case oth.IsSet(): + oth.CPutFn(ctx, batch, kvKey, kvValue, expBytes, traceKV) + case mustValidateOldPKValues: + updateCPutFn(ctx, batch, kvKey, kvValue, expBytes, traceKV, helper, primaryIndexDirs) + default: + // TODO(yuzefovich): in case of multiple column families, + // whenever we locked the primary index during the initial + // scan, we might not have locked the key for a column family + // where all columns had NULL values (because the KV didn't + // exist) and now at least one becomes non-NULL. In this + // scenario we're inserting a new KV with non-locking Put, yet + // we don't have the lock. + // + // However, at the moment we disable the lock eliding + // optimization with multiple column families, so we'll use + // the locking Put because of that. + putFn(ctx, batch, kvKey, kvValue, traceKV, helper, primaryIndexDirs) + } + // Release the kvKey/kvValue references; they're shared across + // calls and must not alias values already handed off to batch. *kvKey = nil - // Prevent future calls to prepareInsertOrUpdateBatch from mutating - // the RawBytes in the kvValue we just added to the batch. Remember - // that we share the kvValue reference across calls to this function. *kvValue = roachpb.Value{} } return rawValueBuf, nil } + +// familyTouchedByUpdate reports whether any column in family appears in +// the updated-column map (i.e., is being inserted or updated this row). +func familyTouchedByUpdate( + family *descpb.ColumnFamilyDescriptor, updatedColIDMapping catalog.TableColMap, +) bool { + for _, colID := range family.ColumnIDs { + if _, ok := updatedColIDMapping.Get(colID); ok { + return true + } + } + return false +} diff --git a/pkg/sql/rowenc/index_encoding.go b/pkg/sql/rowenc/index_encoding.go index efbcb3e27b56..d7054dc2d8cc 100644 --- a/pkg/sql/rowenc/index_encoding.go +++ b/pkg/sql/rowenc/index_encoding.go @@ -1144,6 +1144,200 @@ func encodeVectorIndexKey( return key, err } +// PrimaryIndexFamilyValue is the result of encoding one column family of a +// primary-index row. +// +// Two pieces of state are reported separately: +// +// - Value: the encoded family value, or the zero value when the family +// produced no encoded columns. Callers decide whether absence requires +// a delete (overwrite/CPut paths), a sentinel emission (family 0), or +// no entry at all (pure encoders that omit empty families). +// +// - Skipped: true when the family has no work to do — either the family's +// only column is absent from the input column map, or the column is +// already encoded in the index key and is not composite. Distinguished +// from "Value not present" because callers must NOT issue a delete in +// this case. +type PrimaryIndexFamilyValue struct { + Value roachpb.Value + Skipped bool +} + +// PrimaryIndexEncoder encodes per-family values of a row that uses primary- +// index encoding (the actual primary index, or a covering secondary index +// — anything with PrimaryIndexEncoding). It caches state derived from the +// table descriptor and index (key column set, stored column set, per-family +// sorted column IDs) so that callers can reuse a single encoder across +// every row and every family of an encoding session. +// +// PrimaryIndexEncoder is the canonical place where the per-family value +// encoding rules live. Three independent implementations of these rules +// previously existed in row.prepareInsertOrUpdateBatch, +// row.Deleter.encodeValueForPrimaryIndexFamily, and +// rowenc.EncodePrimaryIndexWithKeyPrefix; all three now route through +// EncodeFamily. +type PrimaryIndexEncoder struct { + desc catalog.TableDescriptor + keyCols catalog.TableColSet + storedCols catalog.TableColSet + // hasStoredCols is false only for indexes encoded with a version + // older than PrimaryIndexWithStoredColumnsVersion (essentially never + // seen in practice). For those, the stored column set is derived + // per-call from the input colMap; see effectiveStoredCols. + hasStoredCols bool + sortedFamilyColIDs map[descpb.FamilyID][]descpb.ColumnID +} + +// MakePrimaryIndexEncoder returns an encoder for the given primary-encoded +// index of desc. The index is typically desc.GetPrimaryIndex(), but it may +// also be a covering secondary index (any index with PrimaryIndexEncoding). +func MakePrimaryIndexEncoder( + desc catalog.TableDescriptor, index catalog.Index, +) PrimaryIndexEncoder { + e := PrimaryIndexEncoder{ + desc: desc, + keyCols: index.CollectKeyColumnIDs(), + } + if index.GetVersion() >= descpb.PrimaryIndexWithStoredColumnsVersion { + if index.Primary() { + e.storedCols = index.CollectPrimaryStoredColumnIDs() + } else { + e.storedCols = index.CollectSecondaryStoredColumnIDs() + } + e.hasStoredCols = true + } + e.sortedFamilyColIDs = make(map[descpb.FamilyID][]descpb.ColumnID, desc.NumFamilies()) + _ = desc.ForeachFamily(func(family *descpb.ColumnFamilyDescriptor) error { + ids := append([]descpb.ColumnID{}, family.ColumnIDs...) + sort.Sort(descpb.ColumnIDs(ids)) + e.sortedFamilyColIDs[family.ID] = ids + return nil + }) + return e +} + +// SkipColumn reports whether the value at colID does not need to be encoded +// in the primary-index value. Composite datums in the key are reported as +// not-skipped; see SkipColumnNotInPrimaryIndexValue for full semantics. +// +// colMap is used only when the primary index is encoded with a version +// older than PrimaryIndexWithStoredColumnsVersion; pass an empty map +// otherwise. +func (e *PrimaryIndexEncoder) SkipColumn( + colID descpb.ColumnID, value tree.Datum, colMap catalog.TableColMap, +) (skip, couldBeComposite bool) { + return SkipColumnNotInPrimaryIndexValue(colID, value, e.keyCols, e.effectiveStoredCols(colMap)) +} + +// SortedFamilyColumnIDs returns the precomputed sorted column IDs for +// family famID, or (nil, false) if no such family exists. +func (e *PrimaryIndexEncoder) SortedFamilyColumnIDs( + famID descpb.FamilyID, +) ([]descpb.ColumnID, bool) { + ids, ok := e.sortedFamilyColIDs[famID] + return ids, ok +} + +// effectiveStoredCols returns the precomputed stored column set when +// available, falling back to a colMap-derived computation for ancient +// primary index versions. +func (e *PrimaryIndexEncoder) effectiveStoredCols(colMap catalog.TableColMap) catalog.TableColSet { + if e.hasStoredCols { + return e.storedCols + } + var all catalog.TableColSet + colMap.ForEach(func(colID descpb.ColumnID, _ int) { + all.Add(colID) + }) + return all.Difference(e.keyCols) +} + +// EncodeFamily encodes the value for one column family of one row of the +// primary index. +// +// colMap maps column ID to position in values. buf is reused as scratch +// for the multi-column tuple encoding; the returned buf may have grown to +// accommodate this call's encoding and should be passed back to the next +// call to amortize allocations. The returned PrimaryIndexFamilyValue's +// Value, when present, owns its own RawBytes (deep-copied via SetTuple or +// MarshalLegacy) and does not alias buf, so the caller is free to mutate +// buf afterward. +// +// The function does NOT compute the family's KV key; callers append the +// family suffix to the index key themselves via keys.MakeFamilyKey. +func (e *PrimaryIndexEncoder) EncodeFamily( + family *descpb.ColumnFamilyDescriptor, + colMap catalog.TableColMap, + values []tree.Datum, + buf []byte, +) (PrimaryIndexFamilyValue, []byte, error) { + storedCols := e.effectiveStoredCols(colMap) + + // Single-default-column family storage optimization: encode the lone + // column directly using legacy value encoding rather than the tuple + // encoding. Family 0 is excluded because its decoder expects a TUPLE + // tag (it carries the row sentinel and composite primary key data). + if len(family.ColumnIDs) == 1 && family.ColumnIDs[0] == family.DefaultColumnID && family.ID != 0 { + idx, ok := colMap.Get(family.DefaultColumnID) + if !ok { + return PrimaryIndexFamilyValue{Skipped: true}, buf, nil + } + skip, couldBeComposite := SkipColumnNotInPrimaryIndexValue( + family.DefaultColumnID, values[idx], e.keyCols, storedCols, + ) + if skip && !couldBeComposite { + return PrimaryIndexFamilyValue{Skipped: true}, buf, nil + } + if skip { + // The column is in the key but is composite, so a previous KV + // may have been written for this family. Return an empty Value; + // overwrite-mode callers must issue a delete to clear it. + return PrimaryIndexFamilyValue{}, buf, nil + } + col, err := catalog.MustFindColumnByID(e.desc, family.DefaultColumnID) + if err != nil { + return PrimaryIndexFamilyValue{}, buf, err + } + marshaled, err := valueside.MarshalLegacy(col.GetType(), values[idx]) + if err != nil { + return PrimaryIndexFamilyValue{}, buf, err + } + return PrimaryIndexFamilyValue{Value: marshaled}, buf, nil + } + + // Multi-column family (or the empty placeholder for family 0). Tuple- + // encode the columns in sorted-ID order with delta-compressed column + // IDs, matching the format expected by valueside decoders. + buf = buf[:0] + var lastColID descpb.ColumnID + for _, colID := range e.sortedFamilyColIDs[family.ID] { + idx, ok := colMap.Get(colID) + if !ok || values[idx] == tree.DNull { + continue + } + if skip, _ := SkipColumnNotInPrimaryIndexValue(colID, values[idx], e.keyCols, storedCols); skip { + continue + } + if lastColID > colID { + return PrimaryIndexFamilyValue{}, buf, + errors.AssertionFailedf("cannot write column id %d after %d", colID, lastColID) + } + colIDDelta := valueside.MakeColumnIDDelta(lastColID, colID) + lastColID = colID + var err error + buf, err = valueside.Encode(buf, colIDDelta, values[idx]) + if err != nil { + return PrimaryIndexFamilyValue{}, buf, err + } + } + var res PrimaryIndexFamilyValue + if family.ID == 0 || len(buf) > 0 { + res.Value.SetTuple(buf) + } + return res, buf, nil +} + // EncodePrimaryIndex constructs the key prefix for the primary index and // delegates the rest of the encoding to EncodePrimaryIndexWithKeyPrefix. func EncodePrimaryIndex( @@ -1158,14 +1352,12 @@ func EncodePrimaryIndex( return EncodePrimaryIndexWithKeyPrefix(tableDesc, index, keyPrefix, colMap, values, includeEmpty) } -// EncodePrimaryIndexWithKeyPrefix constructs a list of k/v pairs for a -// row encoded as a primary index, using the provided key prefix specific to -// that index. This function mirrors the encoding logic in -// prepareInsertOrUpdateBatch in pkg/sql/row/writer.go. It is somewhat -// duplicated here due to the different arguments that -// prepareOrInsertUpdateBatch needs and uses to generate the k/v's for the row -// it inserts. includeEmpty controls whether or not k/v's with empty values -// should be returned. It returns indexEntries in family sorted order. +// EncodePrimaryIndexWithKeyPrefix constructs a list of k/v pairs for a row +// encoded as a primary index, using the provided key prefix specific to +// that index. The per-family value encoding is delegated to +// PrimaryIndexEncoder; this function adds the per-row index key, the +// family key suffixes, and the includeEmpty handling that decides whether +// empty non-zero families show up in the output. func EncodePrimaryIndexWithKeyPrefix( tableDesc catalog.TableDescriptor, index catalog.Index, @@ -1182,74 +1374,41 @@ func EncodePrimaryIndexWithKeyPrefix( return nil, MakeNullPKError(tableDesc, index, colMap, values) } - storedColumns := getStoredColumnsForPrimaryIndex(index, colMap) - keyColumns := index.CollectKeyColumnIDs() - - var entryValue []byte + encoder := MakePrimaryIndexEncoder(tableDesc, index) indexEntries := make([]IndexEntry, 0, tableDesc.NumFamilies()) - var columnsToEncode []ValueEncodedColumn + var entryValueBuf []byte var called bool if err := tableDesc.ForeachFamily(func(family *descpb.ColumnFamilyDescriptor) error { if !called { called = true } else { indexKey = indexKey[:len(indexKey):len(indexKey)] - entryValue = entryValue[:0] - columnsToEncode = columnsToEncode[:0] - } - familyKey := keys.MakeFamilyKey(indexKey, uint32(family.ID)) - // The decoders expect that column family 0 is encoded with a TUPLE value tag, so we - // don't want to use the untagged value encoding. - if len(family.ColumnIDs) == 1 && family.ColumnIDs[0] == family.DefaultColumnID && family.ID != 0 { - // Single column value families which are not stored or key columns can be skipped, - // these may exist temporarily while adding a column. For key columns composite keys may - // be stored at the same time. - if skipColumn, _ := SkipColumnNotInPrimaryIndexValue(family.DefaultColumnID, - values[colMap.GetDefault(family.DefaultColumnID)], - keyColumns, - storedColumns); skipColumn { - return nil - } - datum := findColumnValue(family.DefaultColumnID, colMap, values) - // We want to include this column if its value is non-null or - // we were requested to include all of the columns. - if datum != tree.DNull || includeEmpty { - col, err := catalog.MustFindColumnByID(tableDesc, family.DefaultColumnID) - if err != nil { - return err - } - value, err := valueside.MarshalLegacy(col.GetType(), datum) - if err != nil { - return err - } - indexEntries = append(indexEntries, IndexEntry{Key: familyKey, Value: value, Family: family.ID}) - } - return nil } - - for _, colID := range family.ColumnIDs { - if storedColumns.Contains(colID) { - columnsToEncode = append(columnsToEncode, ValueEncodedColumn{ColID: colID}) - continue - } - if cdatum, ok := values[colMap.GetDefault(colID)].(tree.CompositeDatum); ok { - if cdatum.IsComposite() { - columnsToEncode = append(columnsToEncode, ValueEncodedColumn{ColID: colID, IsComposite: true}) - continue - } - } - } - sort.Sort(ByID(columnsToEncode)) - entryValue, err = writeColumnValues(entryValue, colMap, values, columnsToEncode) + var res PrimaryIndexFamilyValue + res, entryValueBuf, err = encoder.EncodeFamily(family, colMap, values, entryValueBuf) if err != nil { return err } - if family.ID != 0 && len(entryValue) == 0 && !includeEmpty { + if res.Skipped { return nil } - entry := IndexEntry{Key: familyKey, Family: family.ID} - entry.Value.SetTuple(entryValue) - indexEntries = append(indexEntries, entry) + familyKey := keys.MakeFamilyKey(indexKey, uint32(family.ID)) + if !res.Value.IsPresent() { + // Empty family value. Family 0 always emits a TUPLE sentinel + // inside EncodeFamily, so reaching here for family 0 is + // impossible. For non-zero families, only emit when the + // caller asked to see empty entries (e.g. for delete). + if !includeEmpty { + return nil + } + entry := IndexEntry{Key: familyKey, Family: family.ID} + entry.Value.SetTuple(nil) + indexEntries = append(indexEntries, entry) + return nil + } + indexEntries = append(indexEntries, IndexEntry{ + Key: familyKey, Value: res.Value, Family: family.ID, + }) return nil }); err != nil { return nil, err @@ -1264,41 +1423,6 @@ func EncodePrimaryIndexWithKeyPrefix( return indexEntries, nil } -// getStoredColumnsForPrimaryIndex computes the set of columns stored in this -// primary index's value for encoding. Note that EncodePrimaryIndex will utilize -// this set to construct the value, but will augment this with the set of -// key columns which are composite encoded; this is just the set of columns -// stored in the primary index value which are not featured in the index key -// whatsoever. -// -// colMap is expected to include all columns in the table. -func getStoredColumnsForPrimaryIndex( - index catalog.Index, colMap catalog.TableColMap, -) catalog.TableColSet { - - // It should be rare to never that we come across an index which is encoded - // as a primary index but with a version older than this version. - // Nevertheless, for safety, we assume at that version that the stored - // columns set is not populated, and instead we defer to the colMap to - // compute the complete set before subtracting the key columns. - if index.GetVersion() < descpb.PrimaryIndexWithStoredColumnsVersion { - var allColumn catalog.TableColSet - colMap.ForEach(func(colID descpb.ColumnID, _ int) { - allColumn.Add(colID) - }) - return allColumn.Difference(index.CollectKeyColumnIDs()) - } - - // Note that the definition of Primary according to the catalog.Index method - // is that the index is installed as the primary index of the table, not - // that it has a primary index encoding. We must call the appropriate - // method based on this distinction to get the desired set of columns. - if !index.Primary() { - return index.CollectSecondaryStoredColumnIDs() - } - return index.CollectPrimaryStoredColumnIDs() -} - // MakeNullPKError generates an error when the value for a primary key column is // null. func MakeNullPKError( diff --git a/pkg/sql/rowenc/index_encoding_test.go b/pkg/sql/rowenc/index_encoding_test.go index 44c1c92beb4d..a517702625b8 100644 --- a/pkg/sql/rowenc/index_encoding_test.go +++ b/pkg/sql/rowenc/index_encoding_test.go @@ -1482,3 +1482,250 @@ func TestVectorCompositeEncoding(t *testing.T) { require.Equal(t, collatedString, decodedStr) require.Equal(t, 0, len(val)) } + +// TestPrimaryIndexEncoder_EncodeFamily verifies the per-family encoding +// rules implemented by PrimaryIndexEncoder. The cases exercise both +// single-default-column families (which use the legacy direct value +// encoding) and multi-column families (which use the tuple encoding), and +// cover composite primary key columns, NULL values, and the family-0 +// sentinel. +// +// The cross-check at the end compares EncodeFamily output against +// EncodePrimaryIndexWithKeyPrefix on the same tables; both paths must +// produce byte-for-byte identical values. +func TestPrimaryIndexEncoder_EncodeFamily(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + ctx := context.Background() + srv, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{}) + defer srv.Stopper().Stop(ctx) + codec := srv.ApplicationLayer().Codec() + runner := sqlutils.MakeSQLRunner(sqlDB) + + // Helper: load a table descriptor and build a colMap from public + // columns. + loadTable := func(name string) (catalog.TableDescriptor, catalog.TableColMap) { + desc := desctestutils.TestingGetPublicTableDescriptor(kvDB, codec, "defaultdb", name) + var colMap catalog.TableColMap + for _, c := range desc.PublicColumns() { + colMap.Set(c.GetID(), c.Ordinal()) + } + return desc, colMap + } + + // Helper: pull a family by ID. + getFamily := func(desc catalog.TableDescriptor, id descpb.FamilyID) *descpb.ColumnFamilyDescriptor { + for _, f := range desc.GetFamilies() { + if f.ID == id { + famCopy := f + return &famCopy + } + } + t.Fatalf("family %d not found", id) + return nil + } + + t.Run("single default column with value", func(t *testing.T) { + runner.Exec(t, `CREATE TABLE t1 (a INT PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`) + desc, colMap := loadTable("t1") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + values := []tree.Datum{tree.NewDInt(1), tree.NewDInt(42)} + family := getFamily(desc, 1) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.False(t, res.Skipped) + require.True(t, res.Value.IsPresent()) + + // The value should decode to the same int we encoded. + got, err := res.Value.GetInt() + require.NoError(t, err) + require.Equal(t, int64(42), got) + }) + + t.Run("single default column with NULL", func(t *testing.T) { + runner.Exec(t, `CREATE TABLE t2 (a INT PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`) + desc, colMap := loadTable("t2") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + values := []tree.Datum{tree.NewDInt(1), tree.DNull} + family := getFamily(desc, 1) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.False(t, res.Skipped) + // MarshalLegacy of DNull returns an empty Value; callers issue a + // delete when overwriting. + require.False(t, res.Value.IsPresent()) + }) + + t.Run("single default column not in colMap is skipped", func(t *testing.T) { + runner.Exec(t, `CREATE TABLE t3 (a INT PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`) + desc, _ := loadTable("t3") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + // colMap intentionally only contains the PK column. + var colMap catalog.TableColMap + colMap.Set(desc.PublicColumns()[0].GetID(), 0) + values := []tree.Datum{tree.NewDInt(1)} + family := getFamily(desc, 1) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.True(t, res.Skipped) + require.False(t, res.Value.IsPresent()) + }) + + t.Run("multi-column family encodes tuple", func(t *testing.T) { + runner.Exec(t, + `CREATE TABLE t4 (a INT PRIMARY KEY, b INT, c INT, FAMILY pk (a), FAMILY bc (b, c))`, + ) + desc, colMap := loadTable("t4") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + values := []tree.Datum{tree.NewDInt(1), tree.NewDInt(10), tree.NewDInt(20)} + family := getFamily(desc, 1) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.False(t, res.Skipped) + require.True(t, res.Value.IsPresent()) + + // The value tag must be TUPLE. + require.Equal(t, roachpb.ValueType_TUPLE, res.Value.GetTag()) + }) + + t.Run("multi-column family with all NULL values", func(t *testing.T) { + runner.Exec(t, + `CREATE TABLE t5 (a INT PRIMARY KEY, b INT, c INT, FAMILY pk (a), FAMILY bc (b, c))`, + ) + desc, colMap := loadTable("t5") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + values := []tree.Datum{tree.NewDInt(1), tree.DNull, tree.DNull} + family := getFamily(desc, 1) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.False(t, res.Skipped) + // Family != 0 with no encoded columns: Value is empty so the + // caller can decide to delete (overwrite paths) or skip. + require.False(t, res.Value.IsPresent()) + }) + + t.Run("family 0 sentinel always emitted", func(t *testing.T) { + runner.Exec(t, + `CREATE TABLE t6 (a INT PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`, + ) + desc, colMap := loadTable("t6") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + values := []tree.Datum{tree.NewDInt(1), tree.NewDInt(2)} + family := getFamily(desc, 0) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.False(t, res.Skipped) + // Family 0 with only the PK column produces an empty TUPLE — the + // row sentinel — even though no value-side columns were encoded. + require.True(t, res.Value.IsPresent()) + require.Equal(t, roachpb.ValueType_TUPLE, res.Value.GetTag()) + }) + + t.Run("composite primary key column flows into family 0 value", func(t *testing.T) { + runner.Exec(t, + `CREATE TABLE t7 (a DECIMAL PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`, + ) + desc, colMap := loadTable("t7") + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + + // -0 vs 0 — the canonical composite-decimal example. + negZero, err := tree.ParseDDecimal("-0") + require.NoError(t, err) + + values := []tree.Datum{negZero, tree.NewDInt(2)} + family := getFamily(desc, 0) + + res, _, err := enc.EncodeFamily(family, colMap, values, nil) + require.NoError(t, err) + require.True(t, res.Value.IsPresent()) + require.Equal(t, roachpb.ValueType_TUPLE, res.Value.GetTag()) + // Composite encoding of -0 produces non-empty data bytes inside + // the TUPLE; an empty sentinel would have only the header. + require.Greater(t, len(res.Value.RawBytes), 5) + }) + + // Cross-check: EncodeFamily, called over each family, must produce + // the same per-family Values that EncodePrimaryIndexWithKeyPrefix + // produces. This is the single strongest guarantee that the new + // encoder matches the existing pure encoder. + t.Run("matches EncodePrimaryIndexWithKeyPrefix", func(t *testing.T) { + cases := []struct { + ddl string + values []tree.Datum + }{ + { + ddl: `CREATE TABLE x1 (a INT PRIMARY KEY, b INT)`, + values: []tree.Datum{tree.NewDInt(1), tree.NewDInt(2)}, + }, + { + ddl: `CREATE TABLE x2 (a INT PRIMARY KEY, b INT, c INT, FAMILY pk (a), FAMILY bc (b, c))`, + values: []tree.Datum{tree.NewDInt(1), tree.NewDInt(10), tree.NewDInt(20)}, + }, + { + ddl: `CREATE TABLE x3 (a INT PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`, + values: []tree.Datum{tree.NewDInt(1), tree.NewDInt(99)}, + }, + { + ddl: `CREATE TABLE x4 (a INT PRIMARY KEY, b INT, FAMILY pk (a), FAMILY f (b))`, + values: []tree.Datum{tree.NewDInt(1), tree.DNull}, + }, + } + for i, tc := range cases { + t.Run(fmt.Sprintf("case_%d", i), func(t *testing.T) { + runner.Exec(t, tc.ddl) + // Extract the table name from the DDL. + name := fmt.Sprintf("x%d", i+1) + desc, colMap := loadTable(name) + + wantEntries, err := EncodePrimaryIndex( + codec, desc, desc.GetPrimaryIndex(), colMap, tc.values, false, /* includeEmpty */ + ) + require.NoError(t, err) + + // Build a map of family ID -> expected Value bytes from + // the reference encoder. + want := map[descpb.FamilyID][]byte{} + for _, e := range wantEntries { + want[e.Family] = append([]byte(nil), e.Value.RawBytes...) + } + + enc := MakePrimaryIndexEncoder(desc, desc.GetPrimaryIndex()) + var buf []byte + for _, family := range desc.GetFamilies() { + f := family + var res PrimaryIndexFamilyValue + res, buf, err = enc.EncodeFamily(&f, colMap, tc.values, buf) + require.NoError(t, err) + if res.Skipped { + _, hasEntry := want[f.ID] + require.False(t, hasEntry, "family %d Skipped but reference emitted an entry", f.ID) + continue + } + if !res.Value.IsPresent() && f.ID != 0 { + _, hasEntry := want[f.ID] + require.False(t, hasEntry, "family %d empty but reference emitted an entry", f.ID) + continue + } + gotBytes := res.Value.RawBytes + wantBytes, ok := want[f.ID] + require.True(t, ok, "encoder produced family %d but reference omitted it", f.ID) + require.True(t, bytes.Equal(wantBytes, gotBytes), + "family %d: want %x got %x", f.ID, wantBytes, gotBytes) + } + }) + } + }) +}