From 1fb5b63c557926bef495e3a2dcba79c358860012 Mon Sep 17 00:00:00 2001 From: Matt White Date: Wed, 22 Apr 2026 16:42:21 -0700 Subject: [PATCH 1/4] rowenc: add PrimaryIndexEncoder CockroachDB has accumulated four independent implementations of "encode the values for a primary index given a row's datums" - prepareInsertOrUpdateBatch, Deleter.encodeValueForPrimaryIndexFamily, EncodePrimaryIndexWithKeyPrefix, and BatchEncoder.encodePK. The first three operate on []tree.Datum and reimplement the same per-family value encoding skeleton (single-default-column legacy encoding vs. multi- column tuple encoding, family-0 sentinel rules, composite-aware skip checks). This commit adds a single PrimaryIndexEncoder type in rowenc that captures those rules. The encoder caches the table-derived state used by every call (key column set, stored column set, sorted family column IDs), so callers pass only what genuinely varies per row: the family, the column-to-position map, the row's datums, and a reusable scratch buffer. PrimaryIndexFamilyValue distinguishes "Skipped" (no work to do; do not issue a delete) from "Value not present" (no encoded columns; callers in overwrite paths should issue a delete) so the three call sites can share the same per-family decision logic without paying for a delete in the wrong cases. This commit is purely additive. Subsequent commits route the existing encoders through the new helper. Informs #131589. Release note: None Epic: none --- pkg/sql/rowenc/index_encoding.go | 194 ++++++++++++++++++++ pkg/sql/rowenc/index_encoding_test.go | 247 ++++++++++++++++++++++++++ 2 files changed, 441 insertions(+) diff --git a/pkg/sql/rowenc/index_encoding.go b/pkg/sql/rowenc/index_encoding.go index efbcb3e27b56..427bd2b289b2 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( 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) + } + }) + } + }) +} From 589d599703e2644a40f8e0cdf2e1e9341f305468 Mon Sep 17 00:00:00 2001 From: Matt White Date: Wed, 22 Apr 2026 16:43:32 -0700 Subject: [PATCH 2/4] rowenc: route EncodePrimaryIndexWithKeyPrefix through PrimaryIndexEncoder Replace the inlined per-family encoding logic in EncodePrimaryIndexWithKeyPrefix with a call to PrimaryIndexEncoder (added in the previous commit). This removes the duplicated single-default-column / multi-column / family-0 / composite-skip branching, and lets EncodePrimaryIndexWithKeyPrefix focus on what only it cares about: building the index key, applying includeEmpty, and wrapping entries for delete-preserving encoding. getStoredColumnsForPrimaryIndex moves into PrimaryIndexEncoder construction (where it belongs with the other table-derived caches). The unused free function is deleted. Behavior is unchanged in every case observed by the rowenc tests and by the colenc/row equality cross-check (TestEncoderEqualityDatums). There is one rare edge case where behavior is technically improved: when includeEmpty=true and a primary key column lives in a non-zero single-default-column family with a CompositeDatum type whose value is not composite, the new path emits an empty entry where the old path silently skipped. Existing callers either ignore the Value (deleter) or treat an extra empty entry as a no-op delete on a possibly absent key, so the change is harmless in practice and is the more defensible behavior anyway (it lets deleters clean up potentially- stale composite-encoded KVs from prior writes). Informs #131589. Release note: None Epic: none --- pkg/sql/rowenc/index_encoding.go | 126 +++++++------------------------ 1 file changed, 28 insertions(+), 98 deletions(-) diff --git a/pkg/sql/rowenc/index_encoding.go b/pkg/sql/rowenc/index_encoding.go index 427bd2b289b2..d7054dc2d8cc 100644 --- a/pkg/sql/rowenc/index_encoding.go +++ b/pkg/sql/rowenc/index_encoding.go @@ -1352,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, @@ -1376,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 + } + 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 } - entry := IndexEntry{Key: familyKey, Family: family.ID} - entry.Value.SetTuple(entryValue) - indexEntries = append(indexEntries, entry) + indexEntries = append(indexEntries, IndexEntry{ + Key: familyKey, Value: res.Value, Family: family.ID, + }) return nil }); err != nil { return nil, err @@ -1458,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( From 465bbd734ed614a43d0ce7671d1f942f82da6513 Mon Sep 17 00:00:00 2001 From: Matt White Date: Wed, 22 Apr 2026 16:47:44 -0700 Subject: [PATCH 3/4] row: dedup Deleter.encodeValueForPrimaryIndexFamily via PrimaryIndexEncoder RowHelper now embeds a rowenc.PrimaryIndexEncoder, initialized in RowHelper.Init alongside PrimaryIndexKeyPrefix. The three private fields it replaces (primaryIndexKeyCols, primaryIndexValueCols, sortedColumnFamilies) are gone; the encoder owns that state. The two RowHelper public methods that those fields backed - SkipColumnNotInPrimaryIndexValue and SortedColumnFamily - become thin delegations and keep their existing signatures, so colenc and the in-package writer/inserter call sites are unaffected. Deleter.encodeValueForPrimaryIndexFamily collapses to a single delegating call to encoder.EncodeFamily. The "TODO(ssd): Lots of duplication ... rather unfortunate." note that this function was born with goes away too. Note: today's deleter returned an empty Value for the rare "single-default-column composite-but-not-composite-value" case where the writer would have produced an empty Value (and then issued a delete in overwrite mode). Both produce the same caller-visible expected-bytes (nil), so this is a behavior preservation, not a change. Informs #131589. Release note: None Epic: none --- pkg/sql/row/deleter.go | 50 +++++++-------------------------------- pkg/sql/row/helper.go | 53 ++++++++++++++++++++++-------------------- 2 files changed, 37 insertions(+), 66 deletions(-) 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..ab96f657043c 100644 --- a/pkg/sql/row/helper.go +++ b/pkg/sql/row/helper.go @@ -125,11 +125,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 +237,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. @@ -405,6 +414,9 @@ func (rh *RowHelper) encodeSecondaryIndexes( // encodePrimaryIndexValuesToBuf encodes the given values, writing // into the given buffer. +// +// TODO(mw5h): the writer.go family loop will move onto the encoder +// directly in a follow-up commit; this helper will be removed then. func (rh *RowHelper) encodePrimaryIndexValuesToBuf( vals []tree.Datum, valColIDMapping catalog.TableColMap, @@ -440,33 +452,24 @@ func (rh *RowHelper) encodePrimaryIndexValuesToBuf( } // 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 From 225d3974ca0566c1f843338844831f7a73ded646 Mon Sep 17 00:00:00 2001 From: Matt White Date: Wed, 22 Apr 2026 16:52:33 -0700 Subject: [PATCH 4/4] row: dedup prepareInsertOrUpdateBatch via PrimaryIndexEncoder Route prepareInsertOrUpdateBatch's per-family value encoding through RowHelper.PrimaryIndexEncoder().EncodeFamily, the same helper used by the Deleter and by rowenc.EncodePrimaryIndexWithKeyPrefix. The two parallel branches that the function carried for years - one for single-default-column families using legacy direct value encoding, one for multi-column families using tuple encoding - collapse into a single body. Both paths previously reimplemented the family-0 sentinel, the composite-key skip rule, the empty-family delete, and the column-id-delta tuple write; now those rules live in one place. The fetchedCols parameter falls out of the call signature; the encoder looks up types from the table descriptor when it needs them (single-default-column legacy encoding only). The two remaining callers - Inserter.InsertRow and Updater.UpdateRow - pass one fewer argument. The redundant RowHelper.encodePrimaryIndexValuesToBuf method (and the last in-package use of pkg/sql/rowenc/valueside) is removed. Behavior is unchanged for the cases covered by existing tests: TestEncoderEqualityDatums (which compares the row-side write path to the unmodified vector path byte-for-byte), the rowenc/colenc/row unit suites, and the LDR delete tests. Closes #131589. Release note: None Epic: none --- pkg/sql/row/helper.go | 40 ------ pkg/sql/row/inserter.go | 2 +- pkg/sql/row/updater.go | 2 +- pkg/sql/row/writer.go | 293 ++++++++++++++-------------------------- 4 files changed, 106 insertions(+), 231 deletions(-) diff --git a/pkg/sql/row/helper.go b/pkg/sql/row/helper.go index ab96f657043c..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" @@ -412,45 +411,6 @@ func (rh *RowHelper) encodeSecondaryIndexes( return rh.indexEntries, nil } -// encodePrimaryIndexValuesToBuf encodes the given values, writing -// into the given buffer. -// -// TODO(mw5h): the writer.go family loop will move onto the encoder -// directly in a follow-up commit; this helper will be removed then. -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. 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 +}