sql/colfetcher: remove catalog.TableColMap from cFetcherTableArgs

The `catalog.TableColMap` in `cFetcherTableArgs` has been removed.
Information in `IndexFetchSpec` and `cTableInfo` is used instead.

Release note: None

Author: mgartner

pkg/sql/colfetcher/BUILD.bazel

@@ -26,7 +26,6 @@ go_library(
         "//pkg/roachpb",
         "//pkg/settings",
         "//pkg/settings/cluster",
-        "//pkg/sql/catalog",
         "//pkg/sql/catalog/catenumpb",
         "//pkg/sql/catalog/catpb",
         "//pkg/sql/catalog/colinfo",

pkg/sql/colfetcher/cfetcher.go

@@ -178,6 +178,19 @@ func (m colIdxMap) Swap(i, j int) {
 	m.ords[i], m.ords[j] = m.ords[j], m.ords[i]
 }
 
+// Get returns the ordinal for the given column ID, returning -1 if the column
+// ID is not in the map.
+func (m colIdxMap) Get(id descpb.ColumnID) int {
+	// Binary search for the column id.
+	idx := sort.Search(len(m.vals), func(i int) bool {
+		return m.vals[i] >= id
+	})
+	if idx < len(m.vals) && m.vals[idx] == id {
+		return m.ords[idx]
+	}
+	return -1
+}
+
 type cFetcherArgs struct {
 	// memoryLimit determines the maximum memory footprint of the output batch.
 	memoryLimit int64
@@ -387,15 +400,15 @@ func (cf *cFetcher) Init(
 		return errors.AssertionFailedf("unsupported IndexFetchSpec version %d", tableArgs.spec.Version)
 	}
 	table := newCTableInfo()
-	nCols := tableArgs.ColIdxMap.Len()
+	nCols := len(tableArgs.spec.FetchedColumns)
 	if cap(table.orderedColIdxMap.vals) < nCols {
 		table.orderedColIdxMap.vals = make(descpb.ColumnIDs, 0, nCols)
 		table.orderedColIdxMap.ords = make([]int, 0, nCols)
 	}
 	for i := range tableArgs.spec.FetchedColumns {
 		id := tableArgs.spec.FetchedColumns[i].ColumnID
 		table.orderedColIdxMap.vals = append(table.orderedColIdxMap.vals, id)
-		table.orderedColIdxMap.ords = append(table.orderedColIdxMap.ords, tableArgs.ColIdxMap.GetDefault(id))
+		table.orderedColIdxMap.ords = append(table.orderedColIdxMap.ords, i)
 	}
 	sort.Sort(table.orderedColIdxMap)
 	*table = cTableInfo{
@@ -464,10 +477,9 @@ func (cf *cFetcher) Init(
 	needToDecodeDecimalKey := false
 	for i := range fullColumns {
 		col := &fullColumns[i]
-		colIdx, ok := tableArgs.ColIdxMap.Get(col.ColumnID)
-		if ok {
-			//gcassert:bce
-			indexColOrdinals[i] = colIdx
+		//gcassert:bce
+		indexColOrdinals[i] = table.orderedColIdxMap.Get(col.ColumnID)
+		if colIdx := indexColOrdinals[i]; colIdx >= 0 {
 			cf.mustDecodeIndexKey = true
 			needToDecodeDecimalKey = needToDecodeDecimalKey || tableArgs.spec.FetchedColumns[colIdx].Type.Family() == types.DecimalFamily
 			// A composite column might also have a value encoding which must be
@@ -477,9 +489,6 @@ func (cf *cFetcher) Init(
 			} else {
 				table.neededValueColsByIdx.Remove(colIdx)
 			}
-		} else {
-			//gcassert:bce
-			indexColOrdinals[i] = -1
 		}
 	}
 	if needToDecodeDecimalKey && cap(cf.scratch.decoding) < 64 {
@@ -496,8 +505,7 @@ func (cf *cFetcher) Init(
 		suffixCols := table.spec.KeySuffixColumns()
 		for i := range suffixCols {
 			id := suffixCols[i].ColumnID
-			colIdx, ok := tableArgs.ColIdxMap.Get(id)
-			if ok {
+			if colIdx := table.orderedColIdxMap.Get(id); colIdx >= 0 {
 				if suffixCols[i].IsComposite {
 					table.compositeIndexColOrdinals.Add(colIdx)
 					// Note: we account for these composite columns separately: we add
@@ -521,14 +529,8 @@ func (cf *cFetcher) Init(
 		extraValColOrdinals := table.extraValColOrdinals
 		_ = extraValColOrdinals[len(suffixCols)-1]
 		for i := range suffixCols {
-			idx, ok := tableArgs.ColIdxMap.Get(suffixCols[i].ColumnID)
-			if ok {
-				//gcassert:bce
-				extraValColOrdinals[i] = idx
-			} else {
-				//gcassert:bce
-				extraValColOrdinals[i] = -1
-			}
+			//gcassert:bce
+			extraValColOrdinals[i] = table.orderedColIdxMap.Get(suffixCols[i].ColumnID)
 		}
 	}
 
@@ -1150,21 +1152,23 @@ func (cf *cFetcher) processValue(ctx context.Context, familyID descpb.FamilyID)
 			if familyID == 0 {
 				break
 			}
-			// Find the default column ID for the family.
-			var defaultColumnID descpb.ColumnID
+			// Find the default column for the family.
+			var defaultColumnIdx int
+			found := false
 			for _, f := range table.spec.FamilyDefaultColumns {
 				if f.FamilyID == familyID {
-					defaultColumnID = f.DefaultColumnID
+					defaultColumnIdx = table.orderedColIdxMap.Get(f.DefaultColumnID)
+					found = true
 					break
 				}
 			}
-			if defaultColumnID == 0 {
+			if !found {
 				return scrub.WrapError(
 					scrub.IndexKeyDecodingError,
-					errors.AssertionFailedf("single entry value with no default column id"),
+					errors.AssertionFailedf("single entry value with no default column"),
 				)
 			}
-			prettyKey, prettyValue, err = cf.processValueSingle(table, defaultColumnID, prettyKey)
+			prettyKey, prettyValue, err = cf.processValueSingle(table, defaultColumnIdx, prettyKey)
 		}
 		if err != nil {
 			return scrub.WrapError(scrub.IndexValueDecodingError, err)
@@ -1232,35 +1236,35 @@ func (cf *cFetcher) processValue(ctx context.Context, familyID descpb.FamilyID)
 // value in cf.machine.colvecs accordingly.
 // The key is only used for logging.
 func (cf *cFetcher) processValueSingle(
-	table *cTableInfo, colID descpb.ColumnID, prettyKeyPrefix string,
+	table *cTableInfo, idx int, prettyKeyPrefix string,
 ) (prettyKey string, prettyValue string, err error) {
 	prettyKey = prettyKeyPrefix
 
-	if idx, ok := table.ColIdxMap.Get(colID); ok {
-		if cf.traceKV {
-			prettyKey = fmt.Sprintf("%s/%s", prettyKey, table.spec.FetchedColumns[idx].Name)
-		}
-		val := cf.machine.nextKV.Value
-		if !val.IsPresent() {
-			return prettyKey, "", nil
-		}
-		typ := cf.table.spec.FetchedColumns[idx].Type
-		err := colencoding.UnmarshalColumnValueToCol(
-			&table.da, &cf.machine.colvecs, idx, cf.machine.rowIdx, typ, val,
-		)
-		if err != nil {
-			return "", "", err
-		}
-		cf.machine.remainingValueColsByIdx.Remove(idx)
-
-		if cf.traceKV {
-			prettyValue = cf.getDatumAt(idx, cf.machine.rowIdx).String()
-		}
+	if idx < 0 {
+		// No need to unmarshal the column value. Either the column was part of
+		// the index key or it isn't needed.
 		return prettyKey, prettyValue, nil
 	}
 
-	// No need to unmarshal the column value. Either the column was part of
-	// the index key or it isn't needed.
+	if cf.traceKV {
+		prettyKey = fmt.Sprintf("%s/%s", prettyKey, table.spec.FetchedColumns[idx].Name)
+	}
+	val := cf.machine.nextKV.Value
+	if !val.IsPresent() {
+		return prettyKey, "", nil
+	}
+	typ := cf.table.spec.FetchedColumns[idx].Type
+	err = colencoding.UnmarshalColumnValueToCol(
+		&table.da, &cf.machine.colvecs, idx, cf.machine.rowIdx, typ, val,
+	)
+	if err != nil {
+		return "", "", err
+	}
+	cf.machine.remainingValueColsByIdx.Remove(idx)
+
+	if cf.traceKV {
+		prettyValue = cf.getDatumAt(idx, cf.machine.rowIdx).String()
+	}
 	return prettyKey, prettyValue, nil
 }
 

pkg/sql/colfetcher/cfetcher_setup.go

@@ -9,7 +9,6 @@ import (
 	"context"
 	"sync"
 
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/fetchpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
@@ -22,9 +21,6 @@ import (
 // caller) are included in the internal state.
 type cFetcherTableArgs struct {
 	spec fetchpb.IndexFetchSpec
-	// ColIdxMap is a mapping from ColumnID to the ordinal of the corresponding
-	// column within spec.FetchedColumns.
-	ColIdxMap catalog.TableColMap
 	// typs are the types from spec.FetchedColumns.
 	typs []*types.T
 }
@@ -98,9 +94,6 @@ func populateTableArgs(
 		}
 	}
 	args.populateTypes(args.spec.FetchedColumns)
-	for i := range args.spec.FetchedColumns {
-		args.ColIdxMap.Set(args.spec.FetchedColumns[i].ColumnID, i)
-	}
 
 	return args, nil
 }