From c9bbef3935825a228dbd495ab965dba4289c122d Mon Sep 17 00:00:00 2001 From: Marcus Gartner Date: Fri, 19 Aug 2022 08:26:25 -0400 Subject: [PATCH 1/3] opt: return non-pointer from HasMixOfLocalAndRemotePartitions Release note: None --- pkg/sql/opt/partition/locality.go | 7 +++---- pkg/sql/opt/table_meta.go | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/pkg/sql/opt/partition/locality.go b/pkg/sql/opt/partition/locality.go index c86a41ae4ca6..d988bf524bb8 100644 --- a/pkg/sql/opt/partition/locality.go +++ b/pkg/sql/opt/partition/locality.go @@ -145,16 +145,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) { diff --git a/pkg/sql/opt/table_meta.go b/pkg/sql/opt/table_meta.go index 9bc94f1d78ef..ee74899ec483 100644 --- a/pkg/sql/opt/table_meta.go +++ b/pkg/sql/opt/table_meta.go @@ -369,7 +369,7 @@ func (tm *TableMeta) IndexPartitionLocality( if !ok { if localPartitions, ok := partition.HasMixOfLocalAndRemotePartitions(evalCtx, index); ok { - ps = partition.GetSortedPrefixes(index, *localPartitions, evalCtx) + ps = partition.GetSortedPrefixes(index, localPartitions, evalCtx) } tm.AddIndexPartitionLocality(ord, ps) } From afada49cca5d7e2b1891ecf529a5c8eb0e71f1a6 Mon Sep 17 00:00:00 2001 From: Marcus Gartner Date: Fri, 19 Aug 2022 08:41:37 -0400 Subject: [PATCH 2/3] opt: move TableMeta.AddPartialIndexPredicate next to TableMeta.PartialIndexPredicate Release note: None --- pkg/sql/opt/table_meta.go | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/pkg/sql/opt/table_meta.go b/pkg/sql/opt/table_meta.go index ee74899ec483..c74014f1fe1b 100644 --- a/pkg/sql/opt/table_meta.go +++ b/pkg/sql/opt/table_meta.go @@ -318,15 +318,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 @@ -380,6 +371,15 @@ func (tm *TableMeta) IndexPartitionLocality( return ps, ps != nil } +// 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 +} + // PartialIndexPredicate returns the given index's predicate scalar expression, // if the index is a partial index. Returns ok=false if the index is not a // partial index. Panics if the index is a partial index according to the From 06a3d8750904a750518f61178e2aa6ebf96dc5f0 Mon Sep 17 00:00:00 2001 From: Marcus Gartner Date: Fri, 19 Aug 2022 09:14:57 -0400 Subject: [PATCH 3/3] 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 --- pkg/sql/BUILD.bazel | 1 + pkg/sql/constraint.go | 3 +- pkg/sql/opt/BUILD.bazel | 1 - pkg/sql/opt/constraint/constraint.go | 4 +- pkg/sql/opt/constraint/constraint_test.go | 2 +- pkg/sql/opt/constraint/locality.go | 8 +-- pkg/sql/opt/idxconstraint/BUILD.bazel | 1 + .../opt/idxconstraint/index_constraints.go | 2 +- .../idxconstraint/index_constraints_test.go | 5 +- .../opt/invertedidx/inverted_index_expr.go | 2 +- pkg/sql/opt/optbuilder/BUILD.bazel | 1 + pkg/sql/opt/optbuilder/select.go | 15 ++++++ pkg/sql/opt/partition/locality.go | 11 ++-- pkg/sql/opt/table_meta.go | 53 +++++++++---------- pkg/sql/opt/xform/general_funcs.go | 2 +- pkg/sql/opt/xform/join_funcs.go | 17 +++--- pkg/sql/opt/xform/scan_funcs.go | 30 +++++------ pkg/sql/opt/xform/select_funcs.go | 4 +- 18 files changed, 91 insertions(+), 71 deletions(-) diff --git a/pkg/sql/BUILD.bazel b/pkg/sql/BUILD.bazel index f836e9a2ecce..76c59884d39a 100644 --- a/pkg/sql/BUILD.bazel +++ b/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", diff --git a/pkg/sql/constraint.go b/pkg/sql/constraint.go index d22958d5ca16..704afa3d3d6b 100644 --- a/pkg/sql/constraint.go +++ b/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() diff --git a/pkg/sql/opt/BUILD.bazel b/pkg/sql/opt/BUILD.bazel index 3b3d396b243d..af609658d610 100644 --- a/pkg/sql/opt/BUILD.bazel +++ b/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", diff --git a/pkg/sql/opt/constraint/constraint.go b/pkg/sql/opt/constraint/constraint.go index e52b42990cce..e42d52a980a4 100644 --- a/pkg/sql/opt/constraint/constraint.go +++ b/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 diff --git a/pkg/sql/opt/constraint/constraint_test.go b/pkg/sql/opt/constraint/constraint_test.go index 1eb313c58185..b0f34742442c 100644 --- a/pkg/sql/opt/constraint/constraint_test.go +++ b/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) } diff --git a/pkg/sql/opt/constraint/locality.go b/pkg/sql/opt/constraint/locality.go index f3610076494e..f1d988f56b32 100644 --- a/pkg/sql/opt/constraint/locality.go +++ b/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} diff --git a/pkg/sql/opt/idxconstraint/BUILD.bazel b/pkg/sql/opt/idxconstraint/BUILD.bazel index a7a5f2aaab64..288debed2ed5 100644 --- a/pkg/sql/opt/idxconstraint/BUILD.bazel +++ b/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", diff --git a/pkg/sql/opt/idxconstraint/index_constraints.go b/pkg/sql/opt/idxconstraint/index_constraints.go index 091e84921aad..d4451c4b455a 100644 --- a/pkg/sql/opt/idxconstraint/index_constraints.go +++ b/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. diff --git a/pkg/sql/opt/idxconstraint/index_constraints_test.go b/pkg/sql/opt/idxconstraint/index_constraints_test.go index de79f2373255..4859faeeffb9 100644 --- a/pkg/sql/opt/idxconstraint/index_constraints_test.go +++ b/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() diff --git a/pkg/sql/opt/invertedidx/inverted_index_expr.go b/pkg/sql/opt/invertedidx/inverted_index_expr.go index 54e7fcc7fabf..de981e5cdde1 100644 --- a/pkg/sql/opt/invertedidx/inverted_index_expr.go +++ b/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. diff --git a/pkg/sql/opt/optbuilder/BUILD.bazel b/pkg/sql/opt/optbuilder/BUILD.bazel index be7481480c28..27854bdd0353 100644 --- a/pkg/sql/opt/optbuilder/BUILD.bazel +++ b/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", diff --git a/pkg/sql/opt/optbuilder/select.go b/pkg/sql/opt/optbuilder/select.go index b2a2491ef390..9d6060b68ad7 100644 --- a/pkg/sql/opt/optbuilder/select.go +++ b/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) { diff --git a/pkg/sql/opt/partition/locality.go b/pkg/sql/opt/partition/locality.go index d988bf524bb8..bdea23cd81c4 100644 --- a/pkg/sql/opt/partition/locality.go +++ b/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 @@ -173,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()) @@ -214,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 diff --git a/pkg/sql/opt/table_meta.go b/pkg/sql/opt/table_meta.go index c74014f1fe1b..1992c6153973 100644 --- a/pkg/sql/opt/table_meta.go +++ b/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 @@ -345,30 +348,22 @@ 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 } - // 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 + return partition.PrefixSorter{} } // AddPartialIndexPredicate adds a partial index predicate to the table's diff --git a/pkg/sql/opt/xform/general_funcs.go b/pkg/sql/opt/xform/general_funcs.go index d04f3d9d6dd3..e419f5afe7b4 100644 --- a/pkg/sql/opt/xform/general_funcs.go +++ b/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 { diff --git a/pkg/sql/opt/xform/join_funcs.go b/pkg/sql/opt/xform/join_funcs.go index 749daf2386a4..92790099449d 100644 --- a/pkg/sql/opt/xform/join_funcs.go +++ b/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 diff --git a/pkg/sql/opt/xform/scan_funcs.go b/pkg/sql/opt/xform/scan_funcs.go index bcb25f9208b7..e6b8608979a6 100644 --- a/pkg/sql/opt/xform/scan_funcs.go +++ b/pkg/sql/opt/xform/scan_funcs.go @@ -138,12 +138,11 @@ func (c *CustomFuncs) CanMaybeGenerateLocalityOptimizedScan(scanPrivate *memo.Sc // There should be at least two partitions, or we won't be able to // differentiate between local and remote partitions. - // This information is encapsulated in the PrefixSorter. If a PrefixSorter was - // not created for this index, then either all partitions are local, all - // are remote, or the index is not partitioned. + // This information is encapsulated in the PrefixSorter. If a non-empty + // PrefixSorter was not created for this index, then either all partitions + // are local, all are remote, or the index is not partitioned. tabMeta := c.e.mem.Metadata().TableMeta(scanPrivate.Table) - index := tabMeta.Table.Index(scanPrivate.Index) - if _, ok := tabMeta.IndexPartitionLocality(scanPrivate.Index, index, c.e.evalCtx); !ok { + if ps := tabMeta.IndexPartitionLocality(scanPrivate.Index); ps.Empty() { return false } return true @@ -172,14 +171,14 @@ func (c *CustomFuncs) GenerateLocalityOptimizedScan( // 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 - var ok bool - if ps, ok = tabMeta.IndexPartitionLocality(scanPrivate.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(scanPrivate.Index) + if ps.Empty() { return } @@ -191,6 +190,7 @@ func (c *CustomFuncs) GenerateLocalityOptimizedScan( // equivalent to a nil Constraint. idxConstraint := scanPrivate.Constraint if idxConstraint == nil { + var ok bool if idxConstraint, ok = c.buildAllPartitionsConstraint(tabMeta, index, ps, scanPrivate); !ok { return } @@ -297,7 +297,7 @@ func (c *CustomFuncs) GenerateLocalityOptimizedScan( // [/'east' - /'east'/3] [/'east'/5 - /'east']. Adding in the following check // constraint achieves this: CHECK (r IN ('east', 'west', 'central')) func (c *CustomFuncs) buildAllPartitionsConstraint( - tabMeta *opt.TableMeta, index cat.Index, ps *partition.PrefixSorter, sp *memo.ScanPrivate, + tabMeta *opt.TableMeta, index cat.Index, ps partition.PrefixSorter, sp *memo.ScanPrivate, ) (*constraint.Constraint, bool) { var ok bool var remainingFilters memo.FiltersExpr @@ -328,7 +328,7 @@ func (c *CustomFuncs) buildAllPartitionsConstraint( // getLocalSpans returns the indexes of the spans from the given constraint that // target local partitions. func (c *CustomFuncs) getLocalSpans( - scanConstraint *constraint.Constraint, ps *partition.PrefixSorter, + scanConstraint *constraint.Constraint, ps partition.PrefixSorter, ) util.FastIntSet { // Iterate through the spans and determine whether each one matches // with a prefix from a local partition. diff --git a/pkg/sql/opt/xform/select_funcs.go b/pkg/sql/opt/xform/select_funcs.go index 68d75f7fbd8b..89a11f985218 100644 --- a/pkg/sql/opt/xform/select_funcs.go +++ b/pkg/sql/opt/xform/select_funcs.go @@ -173,7 +173,7 @@ func (c *CustomFuncs) MakeCombinedFiltersConstraint( tabMeta *opt.TableMeta, index cat.Index, scanPrivate *memo.ScanPrivate, - ps *partition.PrefixSorter, + ps partition.PrefixSorter, explicitFilters memo.FiltersExpr, optionalFilters memo.FiltersExpr, filterColumns opt.ColSet, @@ -432,7 +432,7 @@ func (c *CustomFuncs) GenerateConstrainedScans( // Create a prefix sorter that describes which index partitions are // local to the gateway region. - prefixSorter, _ := tabMeta.IndexPartitionLocality(index.Ordinal(), index, c.e.evalCtx) + prefixSorter := tabMeta.IndexPartitionLocality(index.Ordinal()) // Build Constraints to scan a subset of the table Spans. if partitionFilters, remainingFilters, combinedConstraint, ok =