From 8397a9eeefbfb12c0a45995a817cd2a2e8b4159c Mon Sep 17 00:00:00 2001 From: Radu Berinde Date: Fri, 18 Feb 2022 12:14:30 -0800 Subject: [PATCH 1/5] sql: remove public key column check when fetching This commit removes a check that fails creation of an IndexFetchSpec if a key column is not public. The only case where this comes up AFAICT is when we hit a duplicate key error and try to fetch the row to get the values for the error message. In this case, even though the column is not public, we know that the row we care about exists in the index, or we wouldn't have hit the error. Release note: None --- pkg/sql/logictest/testdata/logic_test/drop_index | 2 +- pkg/sql/rowenc/index_fetch.go | 6 ------ 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/pkg/sql/logictest/testdata/logic_test/drop_index b/pkg/sql/logictest/testdata/logic_test/drop_index index f7e18659a9b3..f77d4fd6b43d 100644 --- a/pkg/sql/logictest/testdata/logic_test/drop_index +++ b/pkg/sql/logictest/testdata/logic_test/drop_index @@ -326,7 +326,7 @@ DROP INDEX t_secondary CASCADE; ALTER TABLE t DROP COLUMN b; INSERT INTO t SELECT a + 1 FROM t; -statement error pgcode 23505 duplicate key value got decoding error: column \"b\" \(2\) is not public +statement error pgcode 23505 duplicate key value violates unique constraint "t_secondary"\nDETAIL: Key \(b\)=\(0\.0\) already exists UPSERT INTO t SELECT a + 1 FROM t; statement ok diff --git a/pkg/sql/rowenc/index_fetch.go b/pkg/sql/rowenc/index_fetch.go index 96c15b69b7af..08e409930a17 100644 --- a/pkg/sql/rowenc/index_fetch.go +++ b/pkg/sql/rowenc/index_fetch.go @@ -11,8 +11,6 @@ package rowenc import ( - "fmt" - "github.com/cockroachdb/cockroach/pkg/keys" "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" @@ -103,10 +101,6 @@ func InitIndexFetchSpec( } for i := range s.KeyAndSuffixColumns { col := indexCols[i] - if !col.Public() { - // Key columns must be public. - return fmt.Errorf("column %q (%d) is not public", col.GetName(), col.GetID()) - } colID := col.GetID() dir := descpb.IndexDescriptor_ASC // If this is a unique index, the suffix columns are not part of the full From f9240c5054df33f317ca57db34e1f78a3565b723 Mon Sep 17 00:00:00 2001 From: Radu Berinde Date: Fri, 18 Feb 2022 12:08:09 -0800 Subject: [PATCH 2/5] catalog: store IndexFetchSpec key columns in table descriptor We now cache the IndexFetchSpec_KeyAndSuffixColumns in the table descriptor, making it cheaper to populate an IndexFetchSpec. Release note: None --- pkg/sql/catalog/descriptor.go | 4 ++ pkg/sql/catalog/tabledesc/column.go | 54 +++++++++++++++++++++---- pkg/sql/catalog/tabledesc/table_desc.go | 10 +++++ pkg/sql/rowenc/index_fetch.go | 28 +------------ 4 files changed, 61 insertions(+), 35 deletions(-) diff --git a/pkg/sql/catalog/descriptor.go b/pkg/sql/catalog/descriptor.go index 4199898dab41..0c86477e416a 100644 --- a/pkg/sql/catalog/descriptor.go +++ b/pkg/sql/catalog/descriptor.go @@ -509,6 +509,10 @@ type TableDescriptor interface { // stored columns in the specified Index. IndexStoredColumns(idx Index) []Column + // IndexFetchSpecKeyAndSuffixColumns returns information about the key and + // suffix columns, suitable for populating a descpb.IndexFetchSpec. + IndexFetchSpecKeyAndSuffixColumns(idx Index) []descpb.IndexFetchSpec_KeyColumn + // FindColumnWithID returns the first column found whose ID matches the // provided target ID, in the canonical order. // If no column is found then an error is also returned. diff --git a/pkg/sql/catalog/tabledesc/column.go b/pkg/sql/catalog/tabledesc/column.go index 6fb28026784d..782636f15c37 100644 --- a/pkg/sql/catalog/tabledesc/column.go +++ b/pkg/sql/catalog/tabledesc/column.go @@ -265,14 +265,15 @@ type columnCache struct { } type indexColumnCache struct { - all []catalog.Column - allDirs []descpb.IndexDescriptor_Direction - key []catalog.Column - keyDirs []descpb.IndexDescriptor_Direction - stored []catalog.Column - keySuffix []catalog.Column - full []catalog.Column - fullDirs []descpb.IndexDescriptor_Direction + all []catalog.Column + allDirs []descpb.IndexDescriptor_Direction + key []catalog.Column + keyDirs []descpb.IndexDescriptor_Direction + stored []catalog.Column + keySuffix []catalog.Column + full []catalog.Column + fullDirs []descpb.IndexDescriptor_Direction + keyAndSuffix []descpb.IndexFetchSpec_KeyColumn } // newColumnCache returns a fresh fully-populated columnCache struct for the @@ -372,6 +373,43 @@ func makeIndexColumnCache(idx *descpb.IndexDescriptor, all []catalog.Column) (ic } ic.full = ic.all[:nFull] ic.fullDirs = ic.allDirs[:nFull] + + // Populate keyAndSuffix. Note that this method can be called on an incomplete + // (mutable) descriptor (e.g. as part of initializing a new descriptor); this + // code needs to tolerate any descriptor state (like having no key columns, or + // having uninitialized column IDs). + var invertedColumnID descpb.ColumnID + if nKey > 0 && idx.Type == descpb.IndexDescriptor_INVERTED { + invertedColumnID = idx.InvertedColumnID() + } + var compositeIDs catalog.TableColSet + for _, colID := range idx.CompositeColumnIDs { + compositeIDs.Add(colID) + } + ic.keyAndSuffix = make([]descpb.IndexFetchSpec_KeyColumn, nKey+nKeySuffix) + for i := range ic.keyAndSuffix { + col := ic.all[i] + if col == nil { + ic.keyAndSuffix[i].Name = "invalid" + continue + } + colID := col.GetID() + typ := col.GetType() + if colID != 0 && colID == invertedColumnID { + typ = idx.InvertedColumnKeyType() + } + ic.keyAndSuffix[i] = descpb.IndexFetchSpec_KeyColumn{ + IndexFetchSpec_Column: descpb.IndexFetchSpec_Column{ + Name: col.GetName(), + ColumnID: colID, + Type: typ, + IsNonNullable: !col.IsNullable(), + }, + Direction: ic.allDirs[i], + IsComposite: compositeIDs.Contains(colID), + IsInverted: colID == invertedColumnID, + } + } return ic } diff --git a/pkg/sql/catalog/tabledesc/table_desc.go b/pkg/sql/catalog/tabledesc/table_desc.go index 70b080aa668e..c849916c2910 100644 --- a/pkg/sql/catalog/tabledesc/table_desc.go +++ b/pkg/sql/catalog/tabledesc/table_desc.go @@ -492,6 +492,16 @@ func (desc *wrapper) IndexStoredColumns(idx catalog.Index) []catalog.Column { return nil } +// IndexFetchSpecKeyAndSuffixColumns implements the TableDescriptor interface. +func (desc *wrapper) IndexFetchSpecKeyAndSuffixColumns( + idx catalog.Index, +) []descpb.IndexFetchSpec_KeyColumn { + if ic := desc.getExistingOrNewIndexColumnCache(idx); ic != nil { + return ic.keyAndSuffix + } + return nil +} + // getExistingOrNewIndexColumnCache is a convenience method for Index*Columns // methods. func (desc *wrapper) getExistingOrNewIndexColumnCache(idx catalog.Index) *indexColumnCache { diff --git a/pkg/sql/rowenc/index_fetch.go b/pkg/sql/rowenc/index_fetch.go index 08e409930a17..73a8c0ac28b4 100644 --- a/pkg/sql/rowenc/index_fetch.go +++ b/pkg/sql/rowenc/index_fetch.go @@ -32,7 +32,6 @@ func InitIndexFetchSpec( index catalog.Index, fetchColumnIDs []descpb.ColumnID, ) error { - oldKeyAndSuffixCols := s.KeyAndSuffixColumns oldFetchedCols := s.FetchedColumns oldFamilies := s.FamilyDefaultColumns *s = descpb.IndexFetchSpec{ @@ -71,9 +70,7 @@ func InitIndexFetchSpec( } } - indexCols := table.IndexColumns(index) - keyDirs := table.IndexFullColumnDirections(index) - compositeIDs := index.CollectCompositeColumnIDs() + s.KeyAndSuffixColumns = table.IndexFetchSpecKeyAndSuffixColumns(index) var invertedColumnID descpb.ColumnID if index.GetType() == descpb.IndexDescriptor_INVERTED { @@ -93,29 +90,6 @@ func InitIndexFetchSpec( } } - numKeyCols := index.NumKeyColumns() + index.NumKeySuffixColumns() - if cap(oldKeyAndSuffixCols) >= numKeyCols { - s.KeyAndSuffixColumns = oldKeyAndSuffixCols[:numKeyCols] - } else { - s.KeyAndSuffixColumns = make([]descpb.IndexFetchSpec_KeyColumn, numKeyCols) - } - for i := range s.KeyAndSuffixColumns { - col := indexCols[i] - colID := col.GetID() - dir := descpb.IndexDescriptor_ASC - // If this is a unique index, the suffix columns are not part of the full - // index columns and are always ascending. - if i < len(keyDirs) { - dir = keyDirs[i] - } - s.KeyAndSuffixColumns[i] = descpb.IndexFetchSpec_KeyColumn{ - IndexFetchSpec_Column: mkCol(col, colID), - Direction: dir, - IsComposite: compositeIDs.Contains(colID), - IsInverted: colID == invertedColumnID, - } - } - if cap(oldFetchedCols) >= len(fetchColumnIDs) { s.FetchedColumns = oldFetchedCols[:len(fetchColumnIDs)] } else { From c3d1dc1f41d44983fe70cb4729ae7434a7e330a6 Mon Sep 17 00:00:00 2001 From: Radu Berinde Date: Fri, 18 Feb 2022 13:58:21 -0800 Subject: [PATCH 3/5] catalog: cleanup KeysPerRow This change renames KeysPerRow to IndexKeysPerRow and changes the argument to an IndexDescriptor, removing the error case. Release note: None --- pkg/sql/catalog/descriptor.go | 11 ++++++----- pkg/sql/catalog/tabledesc/structured.go | 16 ++++++---------- pkg/sql/catalog/tabledesc/structured_test.go | 7 +++---- pkg/sql/rowenc/index_fetch.go | 9 +++------ pkg/sql/rowexec/joinreader.go | 5 +---- pkg/sql/span/span_splitter.go | 2 +- 6 files changed, 20 insertions(+), 30 deletions(-) diff --git a/pkg/sql/catalog/descriptor.go b/pkg/sql/catalog/descriptor.go index 0c86477e416a..5917ac2b2138 100644 --- a/pkg/sql/catalog/descriptor.go +++ b/pkg/sql/catalog/descriptor.go @@ -345,11 +345,6 @@ type TableDescriptor interface { // 1. Whether the index is a mutation // 2. if so, is it in state DELETE_AND_WRITE_ONLY GetIndexMutationCapabilities(id descpb.IndexID) (isMutation, isWriteOnly bool) - // KeysPerRow returns the maximum number of keys used to encode a row for the - // given index. If a secondary index doesn't store any columns, then it only - // has one k/v pair, but if it stores some columns, it can return up to one - // k/v pair per family in the table, just like a primary index. - KeysPerRow(id descpb.IndexID) (int, error) // AllIndexes returns a slice with all indexes, public and non-public, // in the underlying proto, in their canonical order: @@ -509,6 +504,12 @@ type TableDescriptor interface { // stored columns in the specified Index. IndexStoredColumns(idx Index) []Column + // IndexKeysPerRow returns the maximum number of keys used to encode a row for + // the given index. If a secondary index doesn't store any columns, then it + // only has one k/v pair, but if it stores some columns, it can return up to + // one k/v pair per family in the table, just like a primary index. + IndexKeysPerRow(idx Index) int + // IndexFetchSpecKeyAndSuffixColumns returns information about the key and // suffix columns, suitable for populating a descpb.IndexFetchSpec. IndexFetchSpecKeyAndSuffixColumns(idx Index) []descpb.IndexFetchSpec_KeyColumn diff --git a/pkg/sql/catalog/tabledesc/structured.go b/pkg/sql/catalog/tabledesc/structured.go index 1edc9fe12063..db339aee6941 100644 --- a/pkg/sql/catalog/tabledesc/structured.go +++ b/pkg/sql/catalog/tabledesc/structured.go @@ -105,17 +105,13 @@ func (desc *wrapper) GetParentSchemaID() descpb.ID { return parentSchemaID } -// KeysPerRow implements the TableDescriptor interface. -func (desc *wrapper) KeysPerRow(indexID descpb.IndexID) (int, error) { - if desc.PrimaryIndex.ID == indexID { - return len(desc.Families), nil - } - idx, err := desc.FindIndexWithID(indexID) - if err != nil { - return 0, err +// IndexKeysPerRow implements the TableDescriptor interface. +func (desc *wrapper) IndexKeysPerRow(idx catalog.Index) int { + if desc.PrimaryIndex.ID == idx.GetID() { + return len(desc.Families) } if idx.NumSecondaryStoredColumns() == 0 || len(desc.Families) == 1 { - return 1, nil + return 1 } // Calculate the number of column families used by the secondary index. We // only need to look at the stored columns because column families are only @@ -132,7 +128,7 @@ func (desc *wrapper) KeysPerRow(indexID descpb.IndexID) (int, error) { } } } - return numUsedFamilies, nil + return numUsedFamilies } // BuildIndexName returns an index name that is not equal to any diff --git a/pkg/sql/catalog/tabledesc/structured_test.go b/pkg/sql/catalog/tabledesc/structured_test.go index 5e4cd5454386..1946958e84b9 100644 --- a/pkg/sql/catalog/tabledesc/structured_test.go +++ b/pkg/sql/catalog/tabledesc/structured_test.go @@ -723,10 +723,9 @@ func TestKeysPerRow(t *testing.T) { desc := desctestutils.TestingGetPublicTableDescriptor(db, keys.SystemSQLCodec, "d", tableName) require.NotNil(t, desc) - keys, err := desc.KeysPerRow(test.indexID) - if err != nil { - t.Fatal(err) - } + idx, err := desc.FindIndexWithID(test.indexID) + require.NoError(t, err) + keys := desc.IndexKeysPerRow(idx) if test.expected != keys { t.Errorf("expected %d keys got %d", test.expected, keys) } diff --git a/pkg/sql/rowenc/index_fetch.go b/pkg/sql/rowenc/index_fetch.go index 73a8c0ac28b4..ad096771c3f0 100644 --- a/pkg/sql/rowenc/index_fetch.go +++ b/pkg/sql/rowenc/index_fetch.go @@ -45,13 +45,10 @@ func InitIndexFetchSpec( NumKeySuffixColumns: uint32(index.NumKeySuffixColumns()), } - indexID := index.GetID() - maxKeysPerRow, err := table.KeysPerRow(indexID) - if err != nil { - return err - } + maxKeysPerRow := table.IndexKeysPerRow(index) s.MaxKeysPerRow = uint32(maxKeysPerRow) - s.KeyPrefixLength = uint32(len(MakeIndexKeyPrefix(codec, table.GetID(), indexID))) + // TODO(radu): calculate the length without actually generating a throw-away key. + s.KeyPrefixLength = uint32(len(MakeIndexKeyPrefix(codec, table.GetID(), index.GetID()))) families := table.GetFamilies() for i := range families { diff --git a/pkg/sql/rowexec/joinreader.go b/pkg/sql/rowexec/joinreader.go index 1bbbae2ab115..9f73a3c67db7 100644 --- a/pkg/sql/rowexec/joinreader.go +++ b/pkg/sql/rowexec/joinreader.go @@ -484,10 +484,7 @@ func newJoinReader( ) jr.streamerInfo.unlimitedMemMonitor.Start(flowCtx.EvalCtx.Ctx(), flowCtx.EvalCtx.Mon, mon.BoundAccount{}) jr.streamerInfo.budgetAcc = jr.streamerInfo.unlimitedMemMonitor.MakeBoundAccount() - jr.streamerInfo.maxKeysPerRow, err = jr.desc.KeysPerRow(jr.index.GetID()) - if err != nil { - return nil, err - } + jr.streamerInfo.maxKeysPerRow = jr.desc.IndexKeysPerRow(jr.index) } else { // When not using the Streamer API, we want to limit the batch size hint // to at most half of the workmem limit. Note that it is ok if it is set diff --git a/pkg/sql/span/span_splitter.go b/pkg/sql/span/span_splitter.go index 2e426b594208..072e0feb93bf 100644 --- a/pkg/sql/span/span_splitter.go +++ b/pkg/sql/span/span_splitter.go @@ -72,7 +72,7 @@ func MakeSplitter( // * The index either has just 1 family (so we'll make a GetRequest) or we // need fewer than every column family in the table (otherwise we'd just // make a big ScanRequest). - // TODO(radu): should we be using index.KeysPerRow() instead? + // TODO(radu): should we be using IndexKeysPerRow() instead? numFamilies := len(table.GetFamilies()) if numFamilies > 1 && len(neededFamilies) == numFamilies { return NoopSplitter() From c290a73f7c25de20eed0898cb0abdcddb038f354 Mon Sep 17 00:00:00 2001 From: Radu Berinde Date: Fri, 18 Feb 2022 14:33:49 -0800 Subject: [PATCH 4/5] catalog: cache family default columns in table descriptor We now cache the `IndexFetchSpec_FamilyDefaultColumn`s in the table descriptor, to avoid an allocation when creating an `IndexFetchSpec`. Release note: None --- pkg/sql/catalog/descriptor.go | 4 +++ pkg/sql/catalog/tabledesc/column.go | 37 +++++++++++++++------- pkg/sql/catalog/tabledesc/table_desc.go | 5 +++ pkg/sql/rowenc/index_fetch.go | 41 +++++++++---------------- 4 files changed, 49 insertions(+), 38 deletions(-) diff --git a/pkg/sql/catalog/descriptor.go b/pkg/sql/catalog/descriptor.go index 5917ac2b2138..dbe00f130d2d 100644 --- a/pkg/sql/catalog/descriptor.go +++ b/pkg/sql/catalog/descriptor.go @@ -557,6 +557,10 @@ type TableDescriptor interface { // IDs are unique per table, but not unique globally. GetNextFamilyID() descpb.FamilyID + // FamilyDefaultColumns returns the default column IDs for families with a + // default column. See IndexFetchSpec.FamilyDefaultColumns. + FamilyDefaultColumns() []descpb.IndexFetchSpec_FamilyDefaultColumn + // HasColumnBackfillMutation returns whether the table has any queued column // mutations that require a backfill. HasColumnBackfillMutation() bool diff --git a/pkg/sql/catalog/tabledesc/column.go b/pkg/sql/catalog/tabledesc/column.go index 782636f15c37..ca72c27bba23 100644 --- a/pkg/sql/catalog/tabledesc/column.go +++ b/pkg/sql/catalog/tabledesc/column.go @@ -251,17 +251,18 @@ func (w column) HasGeneratedAsIdentitySequenceOption() bool { // columnCache contains precomputed slices of catalog.Column interfaces. type columnCache struct { - all []catalog.Column - public []catalog.Column - writable []catalog.Column - deletable []catalog.Column - nonDrop []catalog.Column - visible []catalog.Column - accessible []catalog.Column - readable []catalog.Column - withUDTs []catalog.Column - system []catalog.Column - index []indexColumnCache + all []catalog.Column + public []catalog.Column + writable []catalog.Column + deletable []catalog.Column + nonDrop []catalog.Column + visible []catalog.Column + accessible []catalog.Column + readable []catalog.Column + withUDTs []catalog.Column + system []catalog.Column + familyDefaultColumns []descpb.IndexFetchSpec_FamilyDefaultColumn + index []indexColumnCache } type indexColumnCache struct { @@ -337,6 +338,20 @@ func newColumnCache(desc *descpb.TableDescriptor, mutations *mutationCache) *col lazyAllocAppendColumn(&c.withUDTs, col, numDeletable) } } + + // Populate familyDefaultColumns. + for i := range desc.Families { + if f := &desc.Families[i]; f.DefaultColumnID != 0 { + if c.familyDefaultColumns == nil { + c.familyDefaultColumns = make([]descpb.IndexFetchSpec_FamilyDefaultColumn, 0, len(desc.Families)-i) + } + c.familyDefaultColumns = append(c.familyDefaultColumns, descpb.IndexFetchSpec_FamilyDefaultColumn{ + FamilyID: f.ID, + DefaultColumnID: f.DefaultColumnID, + }) + } + } + // Populate the per-index column cache c.index = make([]indexColumnCache, 0, 1+len(desc.Indexes)+len(mutations.indexes)) c.index = append(c.index, makeIndexColumnCache(&desc.PrimaryIndex, c.all)) diff --git a/pkg/sql/catalog/tabledesc/table_desc.go b/pkg/sql/catalog/tabledesc/table_desc.go index c849916c2910..2b239bca8287 100644 --- a/pkg/sql/catalog/tabledesc/table_desc.go +++ b/pkg/sql/catalog/tabledesc/table_desc.go @@ -422,6 +422,11 @@ func (desc *wrapper) SystemColumns() []catalog.Column { return desc.getExistingOrNewColumnCache().system } +// FamilyDefaultColumns implements the TableDescriptor interface. +func (desc *wrapper) FamilyDefaultColumns() []descpb.IndexFetchSpec_FamilyDefaultColumn { + return desc.getExistingOrNewColumnCache().familyDefaultColumns +} + // PublicColumnIDs implements the TableDescriptor interface. func (desc *wrapper) PublicColumnIDs() []descpb.ColumnID { cols := desc.PublicColumns() diff --git a/pkg/sql/rowenc/index_fetch.go b/pkg/sql/rowenc/index_fetch.go index ad096771c3f0..edec233e02a8 100644 --- a/pkg/sql/rowenc/index_fetch.go +++ b/pkg/sql/rowenc/index_fetch.go @@ -33,7 +33,6 @@ func InitIndexFetchSpec( fetchColumnIDs []descpb.ColumnID, ) error { oldFetchedCols := s.FetchedColumns - oldFamilies := s.FamilyDefaultColumns *s = descpb.IndexFetchSpec{ Version: descpb.IndexFetchSpecVersionInitial, TableName: table.GetName(), @@ -50,20 +49,12 @@ func InitIndexFetchSpec( // TODO(radu): calculate the length without actually generating a throw-away key. s.KeyPrefixLength = uint32(len(MakeIndexKeyPrefix(codec, table.GetID(), index.GetID()))) + s.FamilyDefaultColumns = table.FamilyDefaultColumns() + families := table.GetFamilies() for i := range families { - f := &families[i] - if f.DefaultColumnID != 0 { - if s.FamilyDefaultColumns == nil { - s.FamilyDefaultColumns = oldFamilies[:0] - } - s.FamilyDefaultColumns = append(s.FamilyDefaultColumns, descpb.IndexFetchSpec_FamilyDefaultColumn{ - FamilyID: f.ID, - DefaultColumnID: f.DefaultColumnID, - }) - } - if f.ID > s.MaxFamilyID { - s.MaxFamilyID = f.ID + if id := families[i].ID; id > s.MaxFamilyID { + s.MaxFamilyID = id } } @@ -74,19 +65,6 @@ func InitIndexFetchSpec( invertedColumnID = index.InvertedColumnID() } - mkCol := func(col catalog.Column, colID descpb.ColumnID) descpb.IndexFetchSpec_Column { - typ := col.GetType() - if colID == invertedColumnID { - typ = index.InvertedColumnKeyType() - } - return descpb.IndexFetchSpec_Column{ - Name: col.GetName(), - ColumnID: colID, - Type: typ, - IsNonNullable: !col.IsNullable() && col.Public(), - } - } - if cap(oldFetchedCols) >= len(fetchColumnIDs) { s.FetchedColumns = oldFetchedCols[:len(fetchColumnIDs)] } else { @@ -97,7 +75,16 @@ func InitIndexFetchSpec( if err != nil { return err } - s.FetchedColumns[i] = mkCol(col, colID) + typ := col.GetType() + if colID == invertedColumnID { + typ = index.InvertedColumnKeyType() + } + s.FetchedColumns[i] = descpb.IndexFetchSpec_Column{ + Name: col.GetName(), + ColumnID: colID, + Type: typ, + IsNonNullable: !col.IsNullable() && col.Public(), + } } // In test builds, verify that we aren't trying to fetch columns that are not From 9515fc37b7f640db4563fb1da47522adab2c9786 Mon Sep 17 00:00:00 2001 From: Radu Berinde Date: Fri, 18 Feb 2022 14:44:43 -0800 Subject: [PATCH 5/5] rowenc: calculate key prefix length without allocating We now calculate the key prefix length directly instead of creating a throw-away key. Release note: None --- pkg/sql/rowenc/index_fetch.go | 6 ++++-- pkg/util/encoding/encoding.go | 25 +++++++++++++++++++++++++ pkg/util/encoding/encoding_test.go | 11 +++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/pkg/sql/rowenc/index_fetch.go b/pkg/sql/rowenc/index_fetch.go index edec233e02a8..86a54aa01e20 100644 --- a/pkg/sql/rowenc/index_fetch.go +++ b/pkg/sql/rowenc/index_fetch.go @@ -15,6 +15,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/util/buildutil" + "github.com/cockroachdb/cockroach/pkg/util/encoding" "github.com/cockroachdb/errors" ) @@ -46,8 +47,9 @@ func InitIndexFetchSpec( maxKeysPerRow := table.IndexKeysPerRow(index) s.MaxKeysPerRow = uint32(maxKeysPerRow) - // TODO(radu): calculate the length without actually generating a throw-away key. - s.KeyPrefixLength = uint32(len(MakeIndexKeyPrefix(codec, table.GetID(), index.GetID()))) + s.KeyPrefixLength = uint32(len(codec.TenantPrefix()) + + encoding.EncodedLengthUvarintAscending(uint64(s.TableID)) + + encoding.EncodedLengthUvarintAscending(uint64(index.GetID()))) s.FamilyDefaultColumns = table.FamilyDefaultColumns() diff --git a/pkg/util/encoding/encoding.go b/pkg/util/encoding/encoding.go index 634989f3478a..72bda65cef8e 100644 --- a/pkg/util/encoding/encoding.go +++ b/pkg/util/encoding/encoding.go @@ -374,6 +374,31 @@ func EncodeUvarintAscending(b []byte, v uint64) []byte { } } +// EncodedLengthUvarintAscending returns the length of the variable length +// representation, i.e. the number of bytes appended by EncodeUvarintAscending. +func EncodedLengthUvarintAscending(v uint64) int { + switch { + case v <= intSmall: + return 1 + case v <= 0xff: + return 2 + case v <= 0xffff: + return 3 + case v <= 0xffffff: + return 4 + case v <= 0xffffffff: + return 5 + case v <= 0xffffffffff: + return 6 + case v <= 0xffffffffffff: + return 7 + case v <= 0xffffffffffffff: + return 8 + default: + return 9 + } +} + // EncodeUvarintDescending encodes the uint64 value so that it sorts in // reverse order, from largest to smallest. func EncodeUvarintDescending(b []byte, v uint64) []byte { diff --git a/pkg/util/encoding/encoding_test.go b/pkg/util/encoding/encoding_test.go index 926d4f2f3e93..723db2838a85 100644 --- a/pkg/util/encoding/encoding_test.go +++ b/pkg/util/encoding/encoding_test.go @@ -323,6 +323,17 @@ func TestEncodeDecodeUvarint(t *testing.T) { testCustomEncodeUint64(testCases, EncodeUvarintAscending, t) } +func TestEncodedLengthUvarintAscending(t *testing.T) { + for i := 0; i < 100; i++ { + v := rand.Uint64() + exp := len(EncodeUvarintAscending(nil, v)) + actual := EncodedLengthUvarintAscending(v) + if actual != exp { + t.Fatalf("incorrect encoded length for %d: %d (expected %d)", v, actual, exp) + } + } +} + func TestEncodeDecodeUvarintDescending(t *testing.T) { testBasicEncodeDecodeUint64(EncodeUvarintDescending, DecodeUvarintDescending, true, true, true, t) testCases := []testCaseUint64{