Merge #35306

35306: opt: Prefer index with zone constraints that most closely match locality r=andy-kimball a=andy-kimball

Customers can create multiple indexes that are identical, except that they
have different locality constraints. This commit teaches the optimizer to
prefer the index that most closely matches the locality of the gateway node
that is planning the query. This enables scenarios where reference data like
a zip code table can be replicated to different regions, and queries will
use the copy in the same region.

Co-authored-by: Andrew Kimball <[email protected]>

pkg/ccl/logictestccl/testdata/logic_test/zone

@@ -0,0 +1,75 @@
+# LogicTest: 5node-dist-opt
+
+# Ensure that cost-based-optimizer uses an index with zone constraints that most
+# closely matches the gateway's locality.
+
+statement ok
+CREATE TABLE t (
+    k INT PRIMARY KEY,
+    v STRING,
+    INDEX secondary (k) STORING (v)
+);
+
+# ------------------------------------------------------------------------------
+# Put table in dc2 and secondary index in dc1 so that the gateway matches the
+# secondary index rather the primary index.
+# ------------------------------------------------------------------------------
+
+statement ok
+ALTER TABLE t CONFIGURE ZONE USING constraints='[+region=test,+dc=dc2]'
+
+statement ok
+ALTER INDEX t@secondary CONFIGURE ZONE USING constraints='[+region=test,+dc=dc1]'
+
+query TTT
+EXPLAIN SELECT * FROM t WHERE k=10
+----
+scan  ·      ·
+·     table  t@secondary
+·     spans  /10-/11
+
+# ------------------------------------------------------------------------------
+# Swap location of primary and secondary indexes and ensure that primary index
+# is used instead.
+# ------------------------------------------------------------------------------
+
+statement ok
+ALTER TABLE t CONFIGURE ZONE USING constraints='[+region=test,+dc=dc1]'
+
+statement ok
+ALTER INDEX t@secondary CONFIGURE ZONE USING constraints='[+region=test,+dc=dc2]'
+
+query TTT
+EXPLAIN SELECT * FROM t WHERE k=10
+----
+scan  ·      ·
+·     table  t@primary
+·     spans  /10-/10/#
+
+# ------------------------------------------------------------------------------
+# Use PREPARE to make sure that the prepared plan is invalidated when the
+# secondary index's constraints change.
+# ------------------------------------------------------------------------------
+
+statement
+PREPARE p AS SELECT tree, field, description FROM [EXPLAIN SELECT k, v FROM t WHERE k=10]
+
+query TTT
+EXECUTE p
+----
+scan  ·      ·
+·     table  t@primary
+·     spans  /10-/10/#
+
+statement ok
+ALTER TABLE t CONFIGURE ZONE USING constraints='[+region=test,+dc=dc2]'
+
+statement ok
+ALTER INDEX t@secondary CONFIGURE ZONE USING constraints='[+region=test,+dc=dc1]'
+
+query TTT
+EXECUTE p
+----
+scan  ·      ·
+·     table  t@secondary
+·     spans  /10-/11

pkg/config/zone.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -726,3 +727,43 @@ func (z ZoneConfig) subzoneSplits() []roachpb.RKey {
 	}
 	return out
 }
+
+// ReplicaConstraintsCount is part of the cat.Zone interface.
+func (z *ZoneConfig) ReplicaConstraintsCount() int {
+	return len(z.Constraints)
+}
+
+// ReplicaConstraints is part of the cat.Zone interface.
+func (z *ZoneConfig) ReplicaConstraints(i int) cat.ReplicaConstraints {
+	return &z.Constraints[i]
+}
+
+// ReplicaCount is part of the cat.ReplicaConstraints interface.
+func (c *Constraints) ReplicaCount() int32 {
+	return c.NumReplicas
+}
+
+// ConstraintCount is part of the cat.ReplicaConstraints interface.
+func (c *Constraints) ConstraintCount() int {
+	return len(c.Constraints)
+}
+
+// Constraint is part of the cat.ReplicaConstraints interface.
+func (c *Constraints) Constraint(i int) cat.Constraint {
+	return &c.Constraints[i]
+}
+
+// IsRequired is part of the cat.Constraint interface.
+func (c *Constraint) IsRequired() bool {
+	return c.Type == Constraint_REQUIRED
+}
+
+// GetKey is part of the cat.Constraint interface.
+func (c *Constraint) GetKey() string {
+	return c.Key
+}
+
+// GetValue is part of the cat.Constraint interface.
+func (c *Constraint) GetValue() string {
+	return c.Value
+}

pkg/server/server.go

@@ -600,6 +600,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	execCfg = sql.ExecutorConfig{
 		Settings:                s.st,
 		NodeInfo:                nodeInfo,
+		Locality:                s.cfg.Locality,
 		AmbientCtx:              s.cfg.AmbientCtx,
 		DB:                      s.db,
 		Gossip:                  s.gossip,

pkg/sql/conn_executor.go

@@ -1814,6 +1814,7 @@ func (ex *connExecutor) initEvalCtx(ctx context.Context, evalCtx *extendedEvalCo
 			TestingKnobs:     ex.server.cfg.EvalContextTestingKnobs,
 			ClusterID:        ex.server.cfg.ClusterID(),
 			NodeID:           ex.server.cfg.NodeID.Get(),
+			Locality:         ex.server.cfg.Locality,
 			ReCache:          ex.server.reCache,
 			InternalExecutor: ie,
 		},

pkg/sql/exec_util.go

@@ -331,6 +331,7 @@ type nodeStatusGenerator interface {
 type ExecutorConfig struct {
 	Settings *cluster.Settings
 	NodeInfo
+	Locality         roachpb.Locality
 	AmbientCtx       log.AmbientContext
 	DB               *client.DB
 	Gossip           *gossip.Gossip

pkg/sql/logictest/testdata/logic_test/crdb_internal

@@ -290,8 +290,8 @@ node_id  component  field   value
 query ITTTTT colnames
 SELECT node_id, network, regexp_replace(address, '\d+$', '<port>') as address, attrs, locality, regexp_replace(server_version, '^\d+\.\d+(-\d+)?$', '<server_version>') as server_version FROM crdb_internal.gossip_nodes WHERE node_id = 1
 ----
-node_id  network  address           attrs  locality            server_version
-1        tcp      127.0.0.1:<port>  []     {"region": "test"}  <server_version>
+node_id  network  address           attrs  locality                         server_version
+1        tcp      127.0.0.1:<port>  []     {"dc": "dc1", "region": "test"}  <server_version>
 
 query IITBB colnames
 SELECT node_id, epoch, regexp_replace(expiration, '^\d+\.\d+,\d+$', '<timestamp>') as expiration, draining, decommissioning FROM crdb_internal.gossip_liveness WHERE node_id = 1
@@ -303,8 +303,8 @@ query ITTTTTT colnames
 SELECT node_id, network, regexp_replace(address, '\d+$', '<port>') as address, attrs, locality, regexp_replace(server_version, '^\d+\.\d+(-\d+)?$', '<server_version>') as server_version, regexp_replace(go_version, '^go.+$', '<go_version>') as go_version
 FROM crdb_internal.kv_node_status WHERE node_id = 1
 ----
-node_id  network  address           attrs  locality            server_version    go_version
-1        tcp      127.0.0.1:<port>  []     {"region": "test"}  <server_version>  <go_version>
+node_id  network  address           attrs  locality                         server_version    go_version
+1        tcp      127.0.0.1:<port>  []     {"dc": "dc1", "region": "test"}  <server_version>  <go_version>
 
 query IITI colnames
 SELECT node_id, store_id, attrs, used

pkg/sql/opt/cat/index.go

@@ -117,6 +117,16 @@ type Index interface {
 	// of an outbound foreign key relation. Returns false for the second
 	// return value if there is no foreign key reference on this index.
 	ForeignKey() (ForeignKeyReference, bool)
+
+	// Zone returns the zone which constrains placement of the index's range
+	// replicas. If the index was not explicitly assigned to a zone, then it
+	// inherits the zone of its owning table (which in turn inherits from its
+	// owning database or the default zone). In addition, any unspecified zone
+	// information will also be inherited.
+	//
+	// NOTE: This zone always applies to the entire index and never to any
+	// partifular partition of the index.
+	Zone() Zone
 }
 
 // IndexColumn describes a single column that is part of an index definition.

pkg/sql/opt/cat/sequence.go

@@ -30,8 +30,8 @@ type Sequence interface {
 	SequenceName() *tree.TableName
 }
 
-// FormatCatalogSequence nicely formats a catalog sequence using a treeprinter for
+// FormatSequence nicely formats a catalog sequence using a treeprinter for
 // debugging and testing.
-func FormatCatalogSequence(cat Catalog, seq Sequence, tp treeprinter.Node) {
+func FormatSequence(cat Catalog, seq Sequence, tp treeprinter.Node) {
 	tp.Childf("SEQUENCE %s", seq.Name())
 }

pkg/sql/opt/cat/table.go

@@ -202,9 +202,9 @@ func FindTableColumnByName(tab Table, name tree.Name) int {
 	return -1
 }
 
-// FormatCatalogTable nicely formats a catalog table using a treeprinter for
-// debugging and testing.
-func FormatCatalogTable(cat Catalog, tab Table, tp treeprinter.Node) {
+// FormatTable nicely formats a catalog table using a treeprinter for debugging
+// and testing.
+func FormatTable(cat Catalog, tab Table, tp treeprinter.Node) {
 	child := tp.Childf("TABLE %s", tab.Name().TableName)
 
 	var buf bytes.Buffer

pkg/sql/opt/cat/view.go

@@ -40,9 +40,9 @@ type View interface {
 	ColumnName(i int) tree.Name
 }
 
-// FormatCatalogView nicely formats a catalog view using a treeprinter for
-// debugging and testing.
-func FormatCatalogView(view View, tp treeprinter.Node) {
+// FormatView nicely formats a catalog view using a treeprinter for debugging
+// and testing.
+func FormatView(view View, tp treeprinter.Node) {
 	var buf bytes.Buffer
 	if view.ColumnNameCount() > 0 {
 		buf.WriteString(" (")

pkg/sql/opt/cat/zone.go

@@ -0,0 +1,116 @@
+// Copyright 2018 The Cockroach Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package cat
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/util/treeprinter"
+)
+
+// Zone is an interface to zone configuration information used by the optimizer.
+// The optimizer prefers indexes with constraints that best match the locality
+// of the gateway node that plans the query.
+type Zone interface {
+	// ReplicaConstraintsCount returns the number of replica constraint sets that
+	// are part of this zone.
+	ReplicaConstraintsCount() int
+
+	// ReplicaConstraints returns the ith set of replica constraints in the zone,
+	// where i < ReplicaConstraintsCount.
+	ReplicaConstraints(i int) ReplicaConstraints
+}
+
+// ReplicaConstraints is a set of constraints that apply to one or more replicas
+// of a range, restricting which nodes can host that range. For example, if a
+// table range has three replicas, then two of the replicas might be pinned to
+// nodes in one region, whereas the third might be pinned to another region.
+type ReplicaConstraints interface {
+	// ReplicaCount returns the number of replicas that should abide by this set
+	// of constraints. If 0, then the constraints apply to all replicas of the
+	// range (and there can be only one ReplicaConstraints in the Zone).
+	ReplicaCount() int32
+
+	// ConstraintCount returns the number of constraints in the set.
+	ConstraintCount() int
+
+	// Constraint returns the ith constraint in the set, where
+	// i < ConstraintCount.
+	Constraint(i int) Constraint
+}
+
+// Constraint governs placement of range replicas on nodes. A constraint can
+// either be required or prohibited. A required constraint's key/value pair must
+// match one of the tiers of a node's locality for the range to locate there.
+// A prohibited constraint's key/value pair must *not* match any of the tiers of
+// a node's locality for the range to locate there. For example:
+//
+//   +region=east     Range can only be placed on nodes in region=east locality.
+//   -region=west     Range cannot be placed on nodes in region=west locality.
+//
+type Constraint interface {
+	// IsRequired is true if this is a required constraint, or false if this is
+	// a prohibited constraint (signified by initial + or - character).
+	IsRequired() bool
+
+	// GetKey returns the constraint's string key (to left of =).
+	GetKey() string
+
+	// GetValue returns the constraint's string value (to right of =).
+	GetValue() string
+}
+
+// FormatZone nicely formats a catalog zone using a treeprinter for debugging
+// and testing.
+func FormatZone(zone Zone, tp treeprinter.Node) {
+	child := tp.Childf("ZONE")
+	if zone.ReplicaConstraintsCount() > 1 {
+		child = child.Childf("replica constraints")
+	}
+	for i, n := 0, zone.ReplicaConstraintsCount(); i < n; i++ {
+		replConstraint := zone.ReplicaConstraints(i)
+		constraintStr := formatReplicaConstraint(replConstraint)
+		if zone.ReplicaConstraintsCount() > 1 {
+			numReplicas := replConstraint.ReplicaCount()
+			child.Childf("%d replicas: %s", numReplicas, constraintStr)
+		} else {
+			child.Childf("constraints: %s", constraintStr)
+		}
+	}
+}
+
+func formatReplicaConstraint(replConstraint ReplicaConstraints) string {
+	var buf bytes.Buffer
+	buf.WriteRune('[')
+	for i, n := 0, replConstraint.ConstraintCount(); i < n; i++ {
+		constraint := replConstraint.Constraint(i)
+		if i != 0 {
+			buf.WriteRune(',')
+		}
+		if constraint.IsRequired() {
+			buf.WriteRune('+')
+		} else {
+			buf.WriteRune('-')
+		}
+		if constraint.GetKey() != "" {
+			fmt.Fprintf(&buf, "%s=%s", constraint.GetKey(), constraint.GetValue())
+		} else {
+			buf.WriteString(constraint.GetValue())
+		}
+	}
+	buf.WriteRune(']')
+	return buf.String()
+}

pkg/sql/opt/exec/execbuilder/testdata/distsql_agg

@@ -337,15 +337,15 @@ group-by
  ├── grouping columns: b:2
  ├── internal-ordering: +2 opt(1)
  ├── stats: [rows=9.5617925, distinct(2)=9.5617925, null(2)=0]
- ├── cost: 10.7156179
+ ├── cost: 11.1156179
  ├── key: (2)
  ├── fd: (2)-->(3)
  ├── prune: (3)
  ├── scan data2
  │    ├── columns: a:1 b:2
  │    ├── constraint: /1/2: [/1 - /1]
  │    ├── stats: [rows=10, distinct(1)=1, null(1)=0, distinct(2)=9.5617925, null(2)=0]
- │    ├── cost: 10.41
+ │    ├── cost: 10.81
  │    ├── key: (2)
  │    ├── fd: ()-->(1)
  │    ├── ordering: +2 opt(1) [actual: +2]