opt: support locality optimized search for scans with more than 1 row

This commit updates the logic for planning locality optimized search to
allow the optimization the be planned if there are no more than 10,000
keys selected.

The optimization is not yet supported for scans with a hard limit.

Informs cockroachdb#64862

Release justification: Low risk, high benefit change to existing functionality.

Release note (performance improvement): locality optimized search
is now supported for scans that are guaranteed to return 10000 keys
or less. This optimization allows the execution engine to avoid
visiting remote regions if all requested keys are found in the local
region, thus reducing the latency of the query.

pkg/ccl/logictestccl/testdata/logic_test/regional_by_row

@@ -876,6 +876,42 @@ SELECT * FROM [EXPLAIN SELECT * FROM regional_by_row_table WHERE pk = 1] OFFSET
       table: regional_by_row_table@primary
       spans: [/'ap-southeast-2'/1 - /'ap-southeast-2'/1] [/'us-east-1'/1 - /'us-east-1'/1]
 
+# Query with more than one key.
+query T
+SELECT * FROM [EXPLAIN (DISTSQL) SELECT * FROM regional_by_row_table WHERE pk IN (1, 4)] OFFSET 2
+----
+·
+• union all
+│ limit: 2
+│
+├── • scan
+│     missing stats
+│     table: regional_by_row_table@primary
+│     spans: [/'ap-southeast-2'/1 - /'ap-southeast-2'/1] [/'ap-southeast-2'/4 - /'ap-southeast-2'/4]
+│
+└── • scan
+      missing stats
+      table: regional_by_row_table@primary
+      spans: [/'ca-central-1'/1 - /'ca-central-1'/1] [/'ca-central-1'/4 - /'ca-central-1'/4] [/'us-east-1'/1 - /'us-east-1'/1] [/'us-east-1'/4 - /'us-east-1'/4]
+·
+Diagram: https://cockroachdb.github.io/distsqlplan/decode.html#eJykkV-L1DAQwN_9FMM8eRJp0z1B8lTRLi6s27MtKFzLkmuGtdhNapLiHUu_u7R9uKvcCrc-zp_fzG-SE7pfLQrMk23ysYA3sM7SL3CbfL_Zftjs4PWnTV7kX7dXsGywdGiMlu3-7mFvze-9l3ctwbfPSZZA9xNGkjO4vqogXa_zpIAIGWqjaCeP5FDcIseKYWdNTc4ZO6ZOU8NG3aMIGTa66_2YrhjWxhKKE_rGt4QCi3FbRlKRDUJkqMjLpp3GPisWd7Y5SvuADPNOaicgKLEs79-HJQY8CEFqBSsw_gdZhwzT3guIOYsjFq9YfM3id1gNDE3vH6WclwdCwQd2mTi_VDx-Is1n6Rc6R2edH1V7bawiS2qhWQ3PXLUzb00XRMt7ts2x8RCddQhf8m4Zuc5oR3-5nJtcMSR1oPkgZ3pb04019bRmDtOJmxKKnJ-r0Rxs9FSaPvYpzP8Hjv4JrxZwOFTDqz8BAAD__1J1KGk=
+
+statement ok
+SET tracing = on,kv,results; SELECT * FROM regional_by_row_table WHERE pk IN (1, 4); SET tracing = off
+
+# Both rows are found in the local region, so the other regions are not searched.
+query T
+SELECT message FROM [SHOW KV TRACE FOR SESSION] WITH ORDINALITY
+ WHERE message LIKE 'fetched:%' OR message LIKE 'output row%'
+ OR message LIKE 'Scan%'
+ ORDER BY ordinality ASC
+----
+Scan /Table/72/1/"@"/1/0, /Table/72/1/"@"/4/0
+fetched: /regional_by_row_table/primary/'ap-southeast-2'/1/pk2/a/b/j -> /1/2/3/'{"a": "b"}'
+output row: [1 1 2 3 '{"a": "b"}']
+fetched: /regional_by_row_table/primary/'ap-southeast-2'/4/pk2/a/b/j -> /4/5/6/'{"c": "d"}'
+output row: [4 4 5 6 '{"c": "d"}']
+
 # Tests using locality optimized search for lookup joins (including foreign
 # key checks).
 statement ok

pkg/sql/opt/exec/execbuilder/relational.go

@@ -1482,11 +1482,6 @@ func (b *Builder) buildSetOp(set memo.RelExpr) (execPlan, error) {
 		// child.
 		// TODO(rytaft): Store the limit in the expression.
 		hardLimit = uint64(set.Relational().Cardinality.Max)
-		if hardLimit > 1 {
-			panic(errors.AssertionFailedf(
-				"locality optimized search is not yet supported for more than one row at a time",
-			))
-		}
 	}
 
 	ep := execPlan{}

pkg/sql/opt/xform/scan_funcs.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/rowinfra"
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/errors"
 )
@@ -146,11 +147,13 @@ func (c *CustomFuncs) CanMaybeGenerateLocalityOptimizedScan(scanPrivate *memo.Sc
 func (c *CustomFuncs) GenerateLocalityOptimizedScan(
 	grp memo.RelExpr, scanPrivate *memo.ScanPrivate,
 ) {
-	// We can only generate a locality optimized scan if we know there is at
-	// most one row produced by the local spans.
-	// TODO(rytaft): We may be able to expand this to allow any number of rows,
-	// as long as there is a hard upper bound.
-	if !grp.Relational().Cardinality.IsZeroOrOne() {
+	// We can only generate a locality optimized scan if we know there is a hard
+	// upper bound on the number of rows produced by the local spans. We use the
+	// kv batch size as the limit, since it's probably better to use DistSQL once
+	// we're scanning multiple batches.
+	// TODO(rytaft): Revisit this when we have a more accurate cost model for data
+	// distribution.
+	if rowinfra.KeyLimit(grp.Relational().Cardinality.Max) > rowinfra.ProductionKVBatchSize {
 		return
 	}
 

pkg/sql/opt/xform/testdata/rules/scan

@@ -725,23 +725,62 @@ index-join abc_part
            ├── key: ()
            └── fd: ()-->(17-19,21)
 
-# b is not constrained to a single value.
-opt locality=(region=east) expect-not=GenerateLocalityOptimizedScan
+# b is constrained to multiple values.
+opt locality=(region=east) expect=GenerateLocalityOptimizedScan
 SELECT a FROM abc_part WHERE b IN (1, 2)
 ----
 project
  ├── columns: a:3!null
  ├── cardinality: [0 - 2]
  ├── key: (3)
- └── scan abc_part@b_idx
+ └── locality-optimized-search
       ├── columns: a:3!null b:4!null
-      ├── constraint: /1/4
-      │    ├── [/'central'/1 - /'central'/2]
-      │    ├── [/'east'/1 - /'east'/2]
-      │    └── [/'west'/1 - /'west'/2]
+      ├── left columns: a:11 b:12
+      ├── right columns: a:19 b:20
       ├── cardinality: [0 - 2]
       ├── key: (3)
-      └── fd: (3)-->(4), (4)-->(3)
+      ├── fd: (3)-->(4), (4)-->(3)
+      ├── scan abc_part@b_idx
+      │    ├── columns: a:11!null b:12!null
+      │    ├── constraint: /9/12: [/'east'/1 - /'east'/2]
+      │    ├── cardinality: [0 - 2]
+      │    ├── key: (11)
+      │    └── fd: (11)-->(12), (12)-->(11)
+      └── scan abc_part@b_idx
+           ├── columns: a:19!null b:20!null
+           ├── constraint: /17/20
+           │    ├── [/'central'/1 - /'central'/2]
+           │    └── [/'west'/1 - /'west'/2]
+           ├── cardinality: [0 - 4]
+           ├── key: (19)
+           └── fd: (19)-->(20), (20)-->(19)
+
+# b is constrained to more than 100000 values (the kv batch size).
+opt locality=(region=east) expect-not=GenerateLocalityOptimizedScan
+SELECT a FROM abc_part WHERE b >= 0 AND b < 100001
+----
+project
+ ├── columns: a:3!null
+ ├── cardinality: [0 - 100001]
+ ├── key: (3)
+ └── select
+      ├── columns: a:3!null b:4!null
+      ├── cardinality: [0 - 100001]
+      ├── key: (3)
+      ├── fd: (3)-->(4), (4)-->(3)
+      ├── index-join abc_part
+      │    ├── columns: a:3!null b:4
+      │    ├── key: (3)
+      │    ├── fd: (3)-->(4), (4)~~>(3)
+      │    └── scan abc_part@c_idx
+      │         ├── columns: a:3!null
+      │         ├── constraint: /1/2/5
+      │         │    ├── [/'central'/1 - /'central'/3]
+      │         │    ├── [/'east'/1 - /'east'/3]
+      │         │    └── [/'west'/1 - /'west'/3]
+      │         └── key: (3)
+      └── filters
+           └── (b:4 >= 0) AND (b:4 < 100001) [outer=(4), constraints=(/4: [/0 - /100000]; tight)]
 
 # The spans target all remote partitions.
 opt locality=(region=east) expect-not=GenerateLocalityOptimizedScan

pkg/sql/opt_exec_factory.go

@@ -558,11 +558,6 @@ func (ef *execFactory) ConstructStreamingSetOp(
 func (ef *execFactory) ConstructUnionAll(
 	left, right exec.Node, reqOrdering exec.OutputOrdering, hardLimit uint64,
 ) (exec.Node, error) {
-	if hardLimit > 1 {
-		return nil, errors.AssertionFailedf(
-			"locality optimized search is not yet supported for more than one row at a time",
-		)
-	}
 	return ef.planner.newUnionNode(
 		tree.UnionOp,
 		true, /* all */