Merge #86452

86452: opt: reduce index partition locality allocations r=mgartner a=mgartner

#### opt: return non-pointer from HasMixOfLocalAndRemotePartitions

Release note: None

#### opt: move TableMeta.AddPartialIndexPredicate next to TableMeta.PartialIndexPredicate

Release note: None

#### opt: reduce index partition locality allocations

Prior to this commit, a map was allocated in the table metadata to store
`PrefixSorter`s for the indexes of the table. This map was allocated
even if no indexes were partitioned. This commit introduces several
changes that reduce the number of allocations:

  1. The map has been replaced with a slice, which will require only one
     allocation because it will never grow.
  2. The slice is only allocated if there is at least one index with a
     mix of local and remote partitions.
  3. The slice is eagerly added to the table metadata instead of lazily.
     The prefix sorters are accessed in very common rules like
     GenerateConstrainedScans, so there was no benefit in lazily
     creating them.

Release justification: Minor change that fixes micro-benchmark
regressions.

Release note: None


Co-authored-by: Marcus Gartner <[email protected]>

pkg/sql/BUILD.bazel

@@ -378,6 +378,7 @@ go_library(
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/norm",
         "//pkg/sql/opt/optbuilder",
+        "//pkg/sql/opt/partition",
         "//pkg/sql/opt/xform",
         "//pkg/sql/optionalnodeliveness",
         "//pkg/sql/paramparse",

pkg/sql/constraint.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/partition"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
@@ -124,7 +125,7 @@ func (p *planner) ConstrainPrimaryIndexSpanByExpr(
 
 	ic.Init(
 		fe, nil, indexCols, notNullIndexCols, nil,
-		consolidate, evalCtx, &nf, nil,
+		consolidate, evalCtx, &nf, partition.PrefixSorter{},
 	)
 
 	remaining := ic.RemainingFilters()

pkg/sql/opt/BUILD.bazel

@@ -31,7 +31,6 @@ go_library(
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/privilege",
-        "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sem/tree/treebin",
         "//pkg/sql/sem/tree/treecmp",

pkg/sql/opt/constraint/constraint.go

@@ -478,14 +478,14 @@ func (c *Constraint) Combine(evalCtx *eval.Context, other *Constraint) {
 // local partitions will not be consolidated with spans that overlap any remote
 // row ranges. A local row range is one whose leaseholder region preference is
 // the same region as the gateway region.
-func (c *Constraint) ConsolidateSpans(evalCtx *eval.Context, ps *partition.PrefixSorter) {
+func (c *Constraint) ConsolidateSpans(evalCtx *eval.Context, ps partition.PrefixSorter) {
 	keyCtx := KeyContext{Columns: c.Columns, EvalCtx: evalCtx}
 	var result Spans
 
 	if c.Spans.Count() < 1 {
 		return
 	}
-	indexHasLocalAndRemoteParts := ps != nil
+	indexHasLocalAndRemoteParts := !ps.Empty()
 	spanIsLocal, lastSpanIsLocal, localRemoteCrossover := false, false, false
 
 	// Initializations for the first span so we avoid putting a conditional in the

pkg/sql/opt/constraint/constraint_test.go

@@ -516,7 +516,7 @@ func TestConsolidateSpans(t *testing.T) {
 			spans := parseSpans(&evalCtx, tc.s)
 			var c Constraint
 			c.Init(kc, &spans)
-			c.ConsolidateSpans(kc.EvalCtx, nil)
+			c.ConsolidateSpans(kc.EvalCtx, partition.PrefixSorter{})
 			if res := c.Spans.String(); res != tc.e {
 				t.Errorf("expected  %s  got  %s", tc.e, res)
 			}

pkg/sql/opt/constraint/locality.go

@@ -21,7 +21,7 @@ import (
 // compare compares the key prefix in prefixInfo with the span prefix. The key
 // prefix is considered less than the span prefix if it is longer than the
 // span prefix, or if it sorts less according to the Datum.Compare interface.
-func compare(prefixInfo partition.Prefix, span *Span, ps *partition.PrefixSorter) int {
+func compare(prefixInfo partition.Prefix, span *Span, ps partition.PrefixSorter) int {
 	prefix := prefixInfo.Prefix
 	prefixLength := len(prefix)
 	spanPrefixLength := span.Prefix(ps.EvalCtx)
@@ -57,7 +57,7 @@ func compare(prefixInfo partition.Prefix, span *Span, ps *partition.PrefixSorter
 // prefixSearchUpperBound means the same as passing the max upper bound of
 // math.MaxInt32. A zero value for prefixSearchUpperBound means only match on
 // the DEFAULT partition, which has a zero-length prefix.
-func searchPrefixes(span *Span, ps *partition.PrefixSorter, prefixSearchUpperBound int) int {
+func searchPrefixes(span *Span, ps partition.PrefixSorter, prefixSearchUpperBound int) int {
 	if prefixSearchUpperBound < 0 {
 		prefixSearchUpperBound = math.MaxInt32
 	}
@@ -111,7 +111,7 @@ func searchPrefixes(span *Span, ps *partition.PrefixSorter, prefixSearchUpperBou
 // FindMatch finds the Entry in PrefixSorter which matches the span prefix on a
 // prefix subset of its keys, including a zero-length match in the case of the
 // DEFAULT partition.
-func FindMatch(span *Span, ps *partition.PrefixSorter) (match *partition.Prefix, ok bool) {
+func FindMatch(span *Span, ps partition.PrefixSorter) (match *partition.Prefix, ok bool) {
 	index := searchPrefixes(span, ps, math.MaxInt32 /* prefixSearchUpperBound*/)
 	if index == -1 {
 		return nil, false
@@ -123,7 +123,7 @@ func FindMatch(span *Span, ps *partition.PrefixSorter) (match *partition.Prefix,
 // of 1 or less which matches the span prefix, including a zero-length match in
 // the case of the DEFAULT partition.
 func FindMatchOnSingleColumn(
-	datum tree.Datum, ps *partition.PrefixSorter,
+	datum tree.Datum, ps partition.PrefixSorter,
 ) (match *partition.Prefix, ok bool) {
 	sp := &Span{}
 	key := Key{firstVal: datum}

pkg/sql/opt/idxconstraint/BUILD.bazel

@@ -33,6 +33,7 @@ go_test(
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/norm",
         "//pkg/sql/opt/optbuilder",
+        "//pkg/sql/opt/partition",
         "//pkg/sql/opt/testutils",
         "//pkg/sql/parser",
         "//pkg/sql/sem/eval",

pkg/sql/opt/idxconstraint/index_constraints.go

@@ -1054,7 +1054,7 @@ func (ic *Instance) Init(
 	consolidate bool,
 	evalCtx *eval.Context,
 	factory *norm.Factory,
-	ps *partition.PrefixSorter,
+	ps partition.PrefixSorter,
 ) {
 	// This initialization pattern ensures that fields are not unwittingly
 	// reused. Field reuse must be explicit.

pkg/sql/opt/idxconstraint/index_constraints_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/partition"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
@@ -126,7 +127,7 @@ func TestIndexConstraints(t *testing.T) {
 				var ic idxconstraint.Instance
 				ic.Init(
 					filters, optionalFilters, indexCols, sv.NotNullCols(), computedCols,
-					true /* consolidate */, &evalCtx, &f, nil, /* prefixSorter */
+					true /* consolidate */, &evalCtx, &f, partition.PrefixSorter{},
 				)
 				result := ic.Constraint()
 				var buf bytes.Buffer
@@ -242,7 +243,7 @@ func BenchmarkIndexConstraints(b *testing.B) {
 				ic.Init(
 					filters, nil /* optionalFilters */, indexCols, sv.NotNullCols(),
 					nil /* computedCols */, true, /* consolidate */
-					&evalCtx, &f, nil, /* prefixSorter */
+					&evalCtx, &f, partition.PrefixSorter{},
 				)
 				_ = ic.Constraint()
 				_ = ic.RemainingFilters()

pkg/sql/opt/invertedidx/inverted_index_expr.go

@@ -379,7 +379,7 @@ func constrainPrefixColumns(
 ) (constraint *constraint.Constraint, remainingFilters memo.FiltersExpr, ok bool) {
 	tabMeta := factory.Metadata().TableMeta(tabID)
 	prefixColumnCount := index.NonInvertedPrefixColumnCount()
-	ps, _ := tabMeta.IndexPartitionLocality(index.Ordinal(), index, evalCtx)
+	ps := tabMeta.IndexPartitionLocality(index.Ordinal())
 
 	// If this is a single-column inverted index, there are no prefix columns to
 	// constrain.

pkg/sql/opt/optbuilder/BUILD.bazel

@@ -64,6 +64,7 @@ go_library(
         "//pkg/sql/opt/norm",
         "//pkg/sql/opt/optgen/exprgen",
         "//pkg/sql/opt/partialidx",
+        "//pkg/sql/opt/partition",
         "//pkg/sql/opt/props",
         "//pkg/sql/opt/props/physical",
         "//pkg/sql/parser",

pkg/sql/opt/optbuilder/select.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/partition"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
@@ -608,6 +609,7 @@ func (b *Builder) buildScan(
 
 	b.addCheckConstraintsForTable(tabMeta)
 	b.addComputedColsForTable(tabMeta)
+	b.addIndexPartitionLocalitiesForTable(tabMeta)
 
 	outScope.expr = b.factory.ConstructScan(&private)
 
@@ -784,6 +786,19 @@ func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
 	}
 }
 
+// addIndexPartitionLocalitiesForTable caches locality prefix sorters in the
+// table metadata for indexes that have a mix of local and remote partitions.
+func (b *Builder) addIndexPartitionLocalitiesForTable(tabMeta *opt.TableMeta) {
+	tab := tabMeta.Table
+	for indexOrd, n := 0, tab.IndexCount(); indexOrd < n; indexOrd++ {
+		index := tab.Index(indexOrd)
+		if localPartitions, ok := partition.HasMixOfLocalAndRemotePartitions(b.evalCtx, index); ok {
+			ps := partition.GetSortedPrefixes(index, localPartitions, b.evalCtx)
+			tabMeta.AddIndexPartitionLocality(indexOrd, ps)
+		}
+	}
+}
+
 func (b *Builder) buildSequenceSelect(
 	seq cat.Sequence, seqName *tree.TableName, inScope *scope,
 ) (outScope *scope) {

pkg/sql/opt/partition/locality.go

@@ -105,6 +105,11 @@ func (ps PrefixSorter) Swap(i, j int) {
 	ps.Entry[i], ps.Entry[j] = ps.Entry[j], ps.Entry[i]
 }
 
+// Empty returns true if the PrefixSorter contains no prefixes.
+func (ps PrefixSorter) Empty() bool {
+	return len(ps.Entry) == 0
+}
+
 // Slice returns the ith slice of Prefix entries, all having the same
 // partition prefix length, along with the start index of that slice in the
 // main PrefixSorter Entry slice. Slices are sorted by prefix length with those
@@ -145,16 +150,15 @@ func PrefixesToString(prefixes []Prefix) string {
 // determined.
 func HasMixOfLocalAndRemotePartitions(
 	evalCtx *eval.Context, index cat.Index,
-) (localPartitions *util.FastIntSet, ok bool) {
+) (localPartitions util.FastIntSet, ok bool) {
 	if index.PartitionCount() < 2 {
-		return nil, false
+		return util.FastIntSet{}, false
 	}
 	var localRegion string
 	if localRegion, ok = evalCtx.GetLocalRegion(); !ok {
-		return nil, false
+		return util.FastIntSet{}, false
 	}
 	var foundLocal, foundRemote bool
-	localPartitions = &util.FastIntSet{}
 	for i, n := 0, index.PartitionCount(); i < n; i++ {
 		part := index.Partition(i)
 		if IsZoneLocal(part.Zone(), localRegion) {
@@ -174,9 +178,9 @@ func HasMixOfLocalAndRemotePartitions(
 // This is the main function for building a PrefixSorter.
 func GetSortedPrefixes(
 	index cat.Index, localPartitions util.FastIntSet, evalCtx *eval.Context,
-) *PrefixSorter {
+) PrefixSorter {
 	if index == nil || index.PartitionCount() < 2 {
-		return nil
+		return PrefixSorter{}
 	}
 	allPrefixes := make([]Prefix, 0, index.PartitionCount())
 
@@ -215,7 +219,7 @@ func GetSortedPrefixes(
 
 	// The end of the last slice is always the last element.
 	ps.idx = append(ps.idx, len(allPrefixes)-1)
-	return &ps
+	return ps
 }
 
 // isConstraintLocal returns isLocal=true and ok=true if the given constraint is

pkg/sql/opt/table_meta.go

@@ -13,7 +13,6 @@ package opt
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/partition"
-	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/errors"
@@ -166,16 +165,20 @@ type TableMeta struct {
 	// the map.
 	partialIndexPredicates map[cat.IndexOrdinal]ScalarExpr
 
-	// indexPartitionLocalities is a map from an index ordinal on the table to a
-	// *PrefixSorter representing the PARTITION BY LIST values of the index and
-	// whether each of those partitions is region-local with respect to the query
-	// being run. If an index is partitioned BY LIST, and has both local and
-	// remote partitions, it will have an entry in the map. Local partitions are
-	// those where all row ranges they own have a preferred region for leaseholder
-	// nodes the same as the gateway region of the current connection that is
-	// running the query. Remote partitions have at least one row range with a
-	// leaseholder preferred region which is different from the gateway region.
-	indexPartitionLocalities map[cat.IndexOrdinal]*partition.PrefixSorter
+	// indexPartitionLocalities is a slice containing PrefixSorters for each
+	// index that has local and remote partitions. The i-th PrefixSorter in the
+	// slice corresponds to the i-th index in the table.
+	//
+	// The PrefixSorters represent the PARTITION BY LIST values of the index and
+	// whether each of those partitions is region-local with respect to the
+	// query being run. If an index is partitioned BY LIST, and has both local
+	// and remote partitions, it will have an entry in the map. Local partitions
+	// are those where all row ranges they own have a preferred region for
+	// leaseholder nodes the same as the gateway region of the current
+	// connection that is running the query. Remote partitions have at least one
+	// row range with a leaseholder preferred region which is different from the
+	// gateway region.
+	indexPartitionLocalities []partition.PrefixSorter
 
 	// checkConstraintsStats is a map from the current ColumnID statistics based
 	// on CHECK constraint values is based on back to the original ColumnStatistic
@@ -318,15 +321,6 @@ func (tm *TableMeta) ComputedColExpr(id ColumnID) (_ ScalarExpr, ok bool) {
 	return e, ok
 }
 
-// AddPartialIndexPredicate adds a partial index predicate to the table's
-// metadata.
-func (tm *TableMeta) AddPartialIndexPredicate(ord cat.IndexOrdinal, pred ScalarExpr) {
-	if tm.partialIndexPredicates == nil {
-		tm.partialIndexPredicates = make(map[cat.IndexOrdinal]ScalarExpr)
-	}
-	tm.partialIndexPredicates[ord] = pred
-}
-
 // AddCheckConstraintsStats adds a column, ColumnStatistic pair to the
 // checkConstraintsStats map. When the table is duplicated, the mapping from the
 // new check constraint ColumnID back to the original ColumnStatistic is
@@ -354,30 +348,31 @@ func (tm *TableMeta) OrigCheckConstraintsStats(
 
 // AddIndexPartitionLocality adds a PrefixSorter to the table's metadata for the
 // index with IndexOrdinal ord.
-func (tm *TableMeta) AddIndexPartitionLocality(ord cat.IndexOrdinal, ps *partition.PrefixSorter) {
+func (tm *TableMeta) AddIndexPartitionLocality(ord cat.IndexOrdinal, ps partition.PrefixSorter) {
 	if tm.indexPartitionLocalities == nil {
-		tm.indexPartitionLocalities = make(map[cat.IndexOrdinal]*partition.PrefixSorter)
+		tm.indexPartitionLocalities = make([]partition.PrefixSorter, tm.Table.IndexCount())
 	}
 	tm.indexPartitionLocalities[ord] = ps
 }
 
-// IndexPartitionLocality returns the given index's PrefixSorter.
-func (tm *TableMeta) IndexPartitionLocality(
-	ord cat.IndexOrdinal, index cat.Index, evalCtx *eval.Context,
-) (ps *partition.PrefixSorter, ok bool) {
-	ps, ok = tm.indexPartitionLocalities[ord]
-	if !ok {
-		if localPartitions, ok :=
-			partition.HasMixOfLocalAndRemotePartitions(evalCtx, index); ok {
-			ps = partition.GetSortedPrefixes(index, *localPartitions, evalCtx)
-		}
-		tm.AddIndexPartitionLocality(ord, ps)
+// IndexPartitionLocality returns the given index's PrefixSorter. An empty
+// PrefixSorter is returned if the index does not have a mix of local and remote
+// partitions.
+func (tm *TableMeta) IndexPartitionLocality(ord cat.IndexOrdinal) (ps partition.PrefixSorter) {
+	if tm.indexPartitionLocalities != nil {
+		ps := tm.indexPartitionLocalities[ord]
+		return ps
+	}
+	return partition.PrefixSorter{}
+}
+
+// AddPartialIndexPredicate adds a partial index predicate to the table's
+// metadata.
+func (tm *TableMeta) AddPartialIndexPredicate(ord cat.IndexOrdinal, pred ScalarExpr) {
+	if tm.partialIndexPredicates == nil {
+		tm.partialIndexPredicates = make(map[cat.IndexOrdinal]ScalarExpr)
 	}
-	// A nil ps means that the entry in the map for this index indicates that the
-	// index was not partitioned, or the index did not have a mix of local and
-	// remote partitions, so no PrefixSorter is built for it. We return ok=false
-	// to the caller to indicate no PrefixSorter is available for this index.
-	return ps, ps != nil
+	tm.partialIndexPredicates[ord] = pred
 }
 
 // PartialIndexPredicate returns the given index's predicate scalar expression,

pkg/sql/opt/xform/general_funcs.go

@@ -209,7 +209,7 @@ func (c *CustomFuncs) initIdxConstraintForIndex(
 	md := c.e.mem.Metadata()
 	tabMeta := md.TableMeta(tabID)
 	index := tabMeta.Table.Index(indexOrd)
-	ps, _ := tabMeta.IndexPartitionLocality(index.Ordinal(), index, c.e.evalCtx)
+	ps := tabMeta.IndexPartitionLocality(index.Ordinal())
 	columns := make([]opt.OrderingColumn, index.LaxKeyColumnCount())
 	var notNullCols opt.ColSet
 	for i := range columns {

pkg/sql/opt/xform/join_funcs.go

@@ -1289,13 +1289,14 @@ func (c *CustomFuncs) GetLocalityOptimizedLookupJoinExprs(
 
 	// The PrefixSorter has collected all the prefixes from all the different
 	// partitions (remembering which ones came from local partitions), and has
-	// sorted them so that longer prefixes come before shorter prefixes. For each
-	// span in the scanConstraint, we will iterate through the list of prefixes
-	// until we find a match, so ordering them with longer prefixes first ensures
-	// that the correct match is found. The PrefixSorter is only non-nil when this
-	// index has at least one local and one remote partition.
-	var ps *partition.PrefixSorter
-	if ps, ok = tabMeta.IndexPartitionLocality(private.Index, index, c.e.evalCtx); !ok {
+	// sorted them so that longer prefixes come before shorter prefixes. For
+	// each span in the scanConstraint, we will iterate through the list of
+	// prefixes until we find a match, so ordering them with longer prefixes
+	// first ensures that the correct match is found. The PrefixSorter is only
+	// non-empty when this index has at least one local and one remote
+	// partition.
+	ps := tabMeta.IndexPartitionLocality(private.Index)
+	if ps.Empty() {
 		return nil, nil, false
 	}
 
@@ -1365,7 +1366,7 @@ func (c CustomFuncs) getConstPrefixFilter(
 // getLocalValues returns the indexes of the values in the given Datums slice
 // that target local partitions.
 func (c *CustomFuncs) getLocalValues(
-	values tree.Datums, ps *partition.PrefixSorter,
+	values tree.Datums, ps partition.PrefixSorter,
 ) util.FastIntSet {
 	// The PrefixSorter has collected all the prefixes from all the different
 	// partitions (remembering which ones came from local partitions), and has