Merge #72676

72676: catalog: improve how TableColMap handles system columns r=yuzefovich a=yuzefovich

Previously, `TableColMap` would very often incur an allocation when
mapping from system columns are set (because system columns have very
large IDs). This commit teaches `TableColMap` to handle the system
columns separately which allows us to remove those allocations.

Release note: None

Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>

pkg/sql/catalog/BUILD.bazel

@@ -43,17 +43,20 @@ go_test(
  srcs = [
  "dep_test.go",
  "descriptor_test.go",
+ "table_col_map_test.go",
  "table_col_set_test.go",
  ],
  embed = [":catalog"],
  deps = [
+ "//pkg/sql/catalog/colinfo",
  "//pkg/sql/catalog/dbdesc",
  "//pkg/sql/catalog/descpb",
  "//pkg/sql/catalog/schemadesc",
  "//pkg/sql/catalog/tabledesc",
  "//pkg/testutils/buildutil",
  "//pkg/util",
  "//pkg/util/leaktest",
+ "//pkg/util/randutil",
  "@com_github_cockroachdb_redact//:redact",
  "@com_github_stretchr_testify//require",
  "@in_gopkg_yaml_v2//:yaml_v2",

pkg/sql/catalog/catalog.go

@@ -11,6 +11,8 @@
 package catalog
 
 import (
+ "math"
+
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
  "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 )
@@ -103,3 +105,13 @@ func (p ResolvedObjectPrefix) NamePrefix() tree.ObjectNamePrefix {
  }
  return n
 }
+
+// NumSystemColumns defines the number of supported system columns and must be
+// equal to len(colinfo.AllSystemColumnDescs) (enforced in colinfo package to
+// avoid an import cycle).
+const NumSystemColumns = 2
+
+// SmallestSystemColumnColumnID is a descpb.ColumnID with the smallest value
+// among all system columns (enforced in colinfo package to avoid an import
+// cycle).
+const SmallestSystemColumnColumnID = math.MaxUint32 - 1

pkg/sql/catalog/colinfo/system_columns.go

@@ -13,6 +13,7 @@ package colinfo
 import (
  "math"
 
+ "github.com/cockroachdb/cockroach/pkg/sql/catalog"
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
  "github.com/cockroachdb/cockroach/pkg/sql/types"
 )
@@ -44,6 +45,17 @@ const MVCCTimestampColumnID = math.MaxUint32
 // TableOIDColumnID is the ID of the tableoid system column.
 const TableOIDColumnID = MVCCTimestampColumnID - 1
 
+func init() {
+ if len(AllSystemColumnDescs) != catalog.NumSystemColumns {
+ panic("need to update catalog.NumSystemColumns")
+ }
+ for _, desc := range AllSystemColumnDescs {
+ if desc.ID < catalog.SmallestSystemColumnColumnID {
+ panic("need to update catalog.SmallestSystemColumnColumnID")
+ }
+ }
+}
+
 // MVCCTimestampColumnDesc is a column descriptor for the MVCC system column.
 var MVCCTimestampColumnDesc = descpb.ColumnDescriptor{
  Name: MVCCTimestampColumnName,

pkg/sql/catalog/table_col_map.go

@@ -11,42 +11,107 @@
 package catalog
 
 import (
+ "bytes"
+ "fmt"
+
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
  "github.com/cockroachdb/cockroach/pkg/util"
 )
 
 // TableColMap is a map from descpb.ColumnID to int. It is typically used to
 // store a mapping from column id to ordinal position within a row, but can be
 // used for any similar purpose.
+//
+// It stores the mapping for ColumnIDs of the system columns separately since
+// those IDs are very large and incur an allocation in util.FastIntMap all the
+// time.
 type TableColMap struct {
  m util.FastIntMap
+ // systemColMap maps all system columns to their values. Columns here are
+ // in increasing order of their IDs (in other words, since we started giving
+ // out IDs from math.MaxUint32 and are going down, the newer system columns
+ // appear here earlier).
+ systemColMap [NumSystemColumns]int
+ // systemColIsSet indicates whether a value has been set for the
+ // corresponding system column in systemColMap (needed in order to
+ // differentiate between unset 0 and set 0).
+ systemColIsSet [NumSystemColumns]bool
 }
 
 // Set maps a key to the given value.
-func (s *TableColMap) Set(col descpb.ColumnID, val int) { s.m.Set(int(col), val) }
+func (s *TableColMap) Set(col descpb.ColumnID, val int) {
+ if col < SmallestSystemColumnColumnID {
+ s.m.Set(int(col), val)
+ } else {
+ pos := col - SmallestSystemColumnColumnID
+ s.systemColMap[pos] = val
+ s.systemColIsSet[pos] = true
+ }
+}
 
 // Get returns the current value mapped to key, or ok=false if the
 // key is unmapped.
-func (s *TableColMap) Get(col descpb.ColumnID) (val int, ok bool) { return s.m.Get(int(col)) }
+func (s *TableColMap) Get(col descpb.ColumnID) (val int, ok bool) {
+ if col < SmallestSystemColumnColumnID {
+ return s.m.Get(int(col))
+ }
+ pos := col - SmallestSystemColumnColumnID
+ return s.systemColMap[pos], s.systemColIsSet[pos]
+}
 
 // GetDefault returns the current value mapped to key, or 0 if the key is
 // unmapped.
-func (s *TableColMap) GetDefault(col descpb.ColumnID) (val int) { return s.m.GetDefault(int(col)) }
+func (s *TableColMap) GetDefault(col descpb.ColumnID) (val int) {
+ if col < SmallestSystemColumnColumnID {
+ return s.m.GetDefault(int(col))
+ }
+ pos := col - SmallestSystemColumnColumnID
+ return s.systemColMap[pos]
+}
 
 // Len returns the number of keys in the map.
-func (s *TableColMap) Len() (val int) { return s.m.Len() }
+func (s *TableColMap) Len() (val int) {
+ l := s.m.Len()
+ for _, isSet := range s.systemColIsSet {
+ if isSet {
+ l++
+ }
+ }
+ return l
+}
 
 // ForEach calls the given function for each key/value pair in the map (in
 // arbitrary order).
 func (s *TableColMap) ForEach(f func(colID descpb.ColumnID, returnIndex int)) {
  s.m.ForEach(func(k, v int) {
  f(descpb.ColumnID(k), v)
  })
+ for pos, isSet := range s.systemColIsSet {
+ if isSet {
+ id := SmallestSystemColumnColumnID + pos
+ f(descpb.ColumnID(id), s.systemColMap[pos])
+ }
+ }
 }
 
 // String prints out the contents of the map in the following format:
 // map[key1:val1 key2:val2 ...]
 // The keys are in ascending order.
 func (s *TableColMap) String() string {
- return s.m.String()
+ var buf bytes.Buffer
+ buf.WriteString("map[")
+ s.m.ContentsIntoBuffer(&buf)
+ first := buf.Len() == len("map[")
+ for pos, isSet := range s.systemColIsSet {
+ if isSet {
+ if !first {
+ buf.WriteByte(' ')
+ }
+ first = false
+ id := SmallestSystemColumnColumnID + pos
+ fmt.Fprintf(&buf, "%d:%d", id, s.systemColMap[pos])
+ }
+ }
+ buf.WriteByte(']')
+ return buf.String()
 }

pkg/sql/catalog/table_col_map_test.go

@@ -0,0 +1,82 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package catalog_test
+
+import (
+ "math/rand"
+ "sort"
+ "testing"
+
+ "github.com/cockroachdb/cockroach/pkg/sql/catalog"
+ "github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
+ "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+ "github.com/cockroachdb/cockroach/pkg/util"
+ "github.com/cockroachdb/cockroach/pkg/util/randutil"
+ "github.com/stretchr/testify/require"
+)
+
+func TestTableColMap(t *testing.T) {
+ var m catalog.TableColMap
+ var oracle util.FastIntMap
+ rng, _ := randutil.NewTestRand()
+
+ var columnIDs []descpb.ColumnID
+ for i := 0; i < 5; i++ {
+ columnIDs = append(columnIDs, descpb.ColumnID(i))
+ }
+ for _, systemColumnDesc := range colinfo.AllSystemColumnDescs {
+ columnIDs = append(columnIDs, systemColumnDesc.ID)
+ }
+ rand.Shuffle(len(columnIDs), func(i, j int) {
+ columnIDs[i], columnIDs[j] = columnIDs[j], columnIDs[i]
+ })
+
+ // Use each column ID with 50% probability.
+ for i, columnID := range columnIDs {
+ if rng.Float64() < 0.5 {
+ m.Set(columnID, i)
+ oracle.Set(int(columnID), i)
+ }
+ }
+
+ // First, check the length.
+ require.Equal(t, oracle.Len(), m.Len())
+
+ // Check that Get and GetDefault return the same results.
+ for _, columnID := range columnIDs {
+ actual, actualOk := m.Get(columnID)
+ expected, expectedOk := oracle.Get(int(columnID))
+ require.Equal(t, expectedOk, actualOk)
+ if actualOk {
+ require.Equal(t, expected, actual)
+ }
+ actual = m.GetDefault(columnID)
+ expected = oracle.GetDefault(int(columnID))
+ require.Equal(t, expected, actual)
+ }
+
+ // Verify ForEach. We don't bother storing the column IDs here since sorting
+ // them below would be mildly annoying.
+ var actualValues, expectedValues []int
+ m.ForEach(func(_ descpb.ColumnID, returnIndex int) {
+ actualValues = append(actualValues, returnIndex)
+ })
+ oracle.ForEach(func(_ int, val int) {
+ expectedValues = append(expectedValues, val)
+ })
+ // Since the order of iteration is not defined, we have to sort all slices.
+ sort.Ints(actualValues)
+ sort.Ints(expectedValues)
+ require.Equal(t, expectedValues, actualValues)
+
+ // Check that stringification matches too.
+ require.Equal(t, oracle.String(), m.String())
+}

pkg/util/fast_int_map.go

@@ -196,14 +196,12 @@ func (m FastIntMap) ForEach(fn func(key, val int)) {
  }
 }
 
-// String prints out the contents of the map in the following format:
-// map[key1:val1 key2:val2 ...]
+// ContentsIntoBuffer writes the contents of the map into the provided buffer in
+// the following format:
+// key1:val1 key2:val2 ...
 // The keys are in ascending order.
-func (m FastIntMap) String() string {
- var buf bytes.Buffer
- buf.WriteString("map[")
+func (m FastIntMap) ContentsIntoBuffer(buf *bytes.Buffer) {
  first := true
-
  if m.large != nil {
  keys := make([]int, 0, len(m.large))
  for k := range m.large {
@@ -215,7 +213,7 @@ func (m FastIntMap) String() string {
  buf.WriteByte(' ')
  }
  first = false
- fmt.Fprintf(&buf, "%d:%d", k, m.large[k])
+ fmt.Fprintf(buf, "%d:%d", k, m.large[k])
  }
  } else {
  for i := 0; i < numVals; i++ {
@@ -224,10 +222,19 @@ func (m FastIntMap) String() string {
  buf.WriteByte(' ')
  }
  first = false
- fmt.Fprintf(&buf, "%d:%d", i, val)
+ fmt.Fprintf(buf, "%d:%d", i, val)
  }
  }
  }
+}
+
+// String prints out the contents of the map in the following format:
+// map[key1:val1 key2:val2 ...]
+// The keys are in ascending order.
+func (m FastIntMap) String() string {
+ var buf bytes.Buffer
+ buf.WriteString("map[")
+ m.ContentsIntoBuffer(&buf)
  buf.WriteByte(']')
  return buf.String()
 }