opt: Prefer index with zone constraints that most closely match locality

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.

Release note: None

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

@@ -729,13 +729,13 @@ func (z ZoneConfig) subzoneSplits() []roachpb.RKey {
 }
 
 // ReplicaConstraintsCount is part of the cat.Zone interface.
-func (zc *ZoneConfig) ReplicaConstraintsCount() int {
-	return len(zc.Constraints)
+func (z *ZoneConfig) ReplicaConstraintsCount() int {
+	return len(z.Constraints)
 }
 
 // ReplicaConstraints is part of the cat.Zone interface.
-func (zc *ZoneConfig) ReplicaConstraints(i int) cat.ReplicaConstraints {
-	return &zc.Constraints[i]
+func (z *ZoneConfig) ReplicaConstraints(i int) cat.ReplicaConstraints {
+	return &z.Constraints[i]
 }
 
 // ReplicaCount is part of the cat.ReplicaConstraints interface.

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/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]

pkg/sql/opt/exec/execbuilder/testdata/explain

@@ -647,21 +647,21 @@ EXPLAIN (OPT,VERBOSE) SELECT * FROM tc WHERE a = 10 ORDER BY b
 sort
  ├── columns: a:1 b:2
  ├── stats: [rows=9.9, distinct(1)=1, null(1)=0]
- ├── cost: 51.3728708
+ ├── cost: 52.2638708
  ├── fd: ()-->(1)
  ├── ordering: +2 opt(1) [actual: +2]
  ├── prune: (2)
  └── index-join tc
       ├── columns: a:1 b:2
       ├── stats: [rows=9.9, distinct(1)=1, null(1)=0]
-      ├── cost: 50.51
+      ├── cost: 51.401
       ├── fd: ()-->(1)
       ├── prune: (2)
       └── scan tc@c
            ├── columns: a:1 rowid:3
            ├── constraint: /1/3: [/10 - /10]
            ├── stats: [rows=9.9, distinct(1)=1, null(1)=0, distinct(3)=9.9, null(3)=0]
-           ├── cost: 10.306
+           ├── cost: 10.702
            ├── key: (3)
            └── fd: ()-->(1)
 
@@ -671,21 +671,21 @@ EXPLAIN (OPT,TYPES) SELECT * FROM tc WHERE a = 10 ORDER BY b
 sort
  ├── columns: a:1(int!null) b:2(int)
  ├── stats: [rows=9.9, distinct(1)=1, null(1)=0]
- ├── cost: 51.3728708
+ ├── cost: 52.2638708
  ├── fd: ()-->(1)
  ├── ordering: +2 opt(1) [actual: +2]
  ├── prune: (2)
  └── index-join tc
       ├── columns: a:1(int!null) b:2(int)
       ├── stats: [rows=9.9, distinct(1)=1, null(1)=0]
-      ├── cost: 50.51
+      ├── cost: 51.401
       ├── fd: ()-->(1)
       ├── prune: (2)
       └── scan tc@c
            ├── columns: a:1(int!null) rowid:3(int!null)
            ├── constraint: /1/3: [/10 - /10]
            ├── stats: [rows=9.9, distinct(1)=1, null(1)=0, distinct(3)=9.9, null(3)=0]
-           ├── cost: 10.306
+           ├── cost: 10.702
            ├── key: (3)
            └── fd: ()-->(1)
 
@@ -707,24 +707,24 @@ EXPLAIN (OPT, VERBOSE) SELECT * FROM tc WHERE a + 2 * b > 1 ORDER BY a*b
 sort
  ├── columns: a:1 b:2  [hidden: column4:4]
  ├── stats: [rows=333.333333]
- ├── cost: 1129.24548
+ ├── cost: 1179.24548
  ├── fd: (1,2)-->(4)
  ├── ordering: +4
  ├── prune: (1,2,4)
  └── project
       ├── columns: column4:4 a:1 b:2
       ├── stats: [rows=333.333333]
-      ├── cost: 1066.69667
+      ├── cost: 1116.69667
       ├── fd: (1,2)-->(4)
       ├── prune: (1,2,4)
       ├── select
       │    ├── columns: a:1 b:2
       │    ├── stats: [rows=333.333333]
-      │    ├── cost: 1060.02
+      │    ├── cost: 1110.02
       │    ├── scan tc
       │    │    ├── columns: a:1 b:2
       │    │    ├── stats: [rows=1000]
-      │    │    ├── cost: 1050.01
+      │    │    ├── cost: 1100.01
       │    │    └── prune: (1,2)
       │    └── filters
       │         └── (a + (b * 2)) > 1 [outer=(1,2)]
@@ -737,24 +737,24 @@ EXPLAIN (OPT, TYPES) SELECT * FROM tc WHERE a + 2 * b > 1 ORDER BY a*b
 sort
  ├── columns: a:1(int) b:2(int)  [hidden: column4:4(int)]
  ├── stats: [rows=333.333333]
- ├── cost: 1129.24548
+ ├── cost: 1179.24548
  ├── fd: (1,2)-->(4)
  ├── ordering: +4
  ├── prune: (1,2,4)
  └── project
       ├── columns: column4:4(int) a:1(int) b:2(int)
       ├── stats: [rows=333.333333]
-      ├── cost: 1066.69667
+      ├── cost: 1116.69667
       ├── fd: (1,2)-->(4)
       ├── prune: (1,2,4)
       ├── select
       │    ├── columns: a:1(int) b:2(int)
       │    ├── stats: [rows=333.333333]
-      │    ├── cost: 1060.02
+      │    ├── cost: 1110.02
       │    ├── scan tc
       │    │    ├── columns: a:1(int) b:2(int)
       │    │    ├── stats: [rows=1000]
-      │    │    ├── cost: 1050.01
+      │    │    ├── cost: 1100.01
       │    │    └── prune: (1,2)
       │    └── filters
       │         └── gt [type=bool, outer=(1,2)]

pkg/sql/opt/testutils/opttester/opt_tester.go

@@ -24,6 +24,7 @@ import (
 	"testing"
 	"text/tabwriter"
 
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
@@ -122,6 +123,15 @@ type Flags struct {
 	// ReorderJoinsLimit is the maximum number of joins in a query which the optimizer
 	// should attempt to reorder.
 	JoinLimit int
+
+	// Locality specifies the location of the planning node as a set of user-
+	// defined key/value pairs, ordered from most inclusive to least inclusive.
+	// If there are no tiers, then the node's location is not known. Examples:
+	//
+	//   [region=eu]
+	//   [region=us,dc=east]
+	//
+	Locality roachpb.Locality
 }
 
 // New constructs a new instance of the OptTester for the given SQL statement.
@@ -223,6 +233,10 @@ func New(catalog cat.Catalog, sql string) *OptTester {
 //    expression in the query tree for the purpose of creating alternate query
 //    plans in the optimizer.
 //
+//  - locality: used to set the locality of the node that plans the query. This
+//    can affect costing when there are multiple possible indexes to choose
+//    from, each in different localities.
+//
 func (ot *OptTester) RunCommand(tb testing.TB, d *datadriven.TestData) string {
 	// Allow testcases to override the flags.
 	for _, a := range d.CmdArgs {
@@ -240,6 +254,7 @@ func (ot *OptTester) RunCommand(tb testing.TB, d *datadriven.TestData) string {
 
 	ot.Flags.Verbose = testing.Verbose()
 	ot.evalCtx.TestingKnobs.OptimizerCostPerturbation = ot.Flags.PerturbCost
+	ot.evalCtx.Locality = ot.Flags.Locality
 
 	switch d.Cmd {
 	case "exec-ddl":
@@ -509,6 +524,14 @@ func (f *Flags) Set(arg datadriven.CmdArg) error {
 			return err
 		}
 
+	case "locality":
+		// Recombine multiple arguments, separated by commas.
+		locality := strings.Join(arg.Vals, ",")
+		err := f.Locality.Set(locality)
+		if err != nil {
+			return err
+		}
+
 	default:
 		return fmt.Errorf("unknown argument: %s", arg.Key)
 	}

pkg/sql/opt/testutils/testcat/set_zone_config.go

@@ -24,10 +24,6 @@ import (
 
 // SetZoneConfig is a partial implementation of the ALTER TABLE ... CONFIGURE
 // ZONE USING statement.
-//
-// Supported commands:
-//  - INJECT STATISTICS: imports table statistics from a JSON object.
-//
 func (tc *Catalog) SetZoneConfig(stmt *tree.SetZoneConfig) *config.ZoneConfig {
 	// Update the table name to include catalog and schema if not provided.
 	tabName := stmt.TableOrIndex.Table