opt: use optimizer_span_limit in CanMaybeGenerateLocalityOptimizedScan

Use the optimizer_span_limit session setting to further limit the
maximum number of spans allowed in CanMaybeGenerateLocalityOptimizedScan.
When set, the span limit replaces the default 10000 span cap if it is
lower, preventing locality-optimized scans from being generated for
queries that produce too many constraint spans.

Release note: None

Author: michae2

pkg/sql/opt/xform/scan_funcs.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/rowinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/intsets"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 )
 
@@ -136,7 +137,15 @@ func (c *CustomFuncs) CanMaybeGenerateLocalityOptimizedScan(scanPrivate *memo.Sc
 
 		// Don't apply the rule if there are too many spans, since the rule code is
 		// O(# spans * # prefixes * # datums per prefix).
-		if scanPrivate.Constraint.Spans.Count() > 10000 {
+		maxSpans := 10000
+		if spanLimit := int(c.e.evalCtx.SessionData().OptimizerSpanLimit); spanLimit > 0 && spanLimit < maxSpans {
+			maxSpans = spanLimit
+		}
+		if scanPrivate.Constraint.Spans.Count() > maxSpans {
+			log.VEventf(c.e.ctx, 2,
+				"optimizer_span_limit: locality-optimized scan has %d spans, exceeds limit of %d",
+				scanPrivate.Constraint.Spans.Count(), maxSpans,
+			)
 			return false
 		}
 	}
@@ -417,6 +426,7 @@ func (c *CustomFuncs) getLocalAndRemoteFilters(
 // in the local (left) branch of the union all should be
 // [/'east' - /'east'/3] [/'east'/5 - /'east']. Adding in the following check
 // constraint achieves this: CHECK (r IN ('east', 'west', 'central'))
+// TODO(michae2): Pass optimizer_span_limit to MakeCombinedFiltersConstraint.
 func (c *CustomFuncs) buildAllPartitionsConstraint(
 	tabMeta *opt.TableMeta, index cat.Index, ps partition.PrefixSorter, sp *memo.ScanPrivate,
 ) (*constraint.Constraint, bool) {

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

@@ -1318,3 +1318,102 @@ locality-optimized-search
       ├── limit: 3
       ├── flags: force-index=def_part_pkey
       └── key: (14)
+
+# --------------------------------------------------
+# optimizer_span_limit with GenerateLocalityOptimizedScan
+# --------------------------------------------------
+
+# b = 1 produces 3 spans (one per partition). By default, the locality-optimized
+# scan fires.
+opt locality=(region=east) expect=GenerateLocalityOptimizedScan
+SELECT a FROM abc_part WHERE b = 1
+----
+project
+ ├── columns: a:3!null
+ ├── cardinality: [0 - 1]
+ ├── key: ()
+ ├── fd: ()-->(3)
+ ├── distribution: east
+ └── locality-optimized-search
+      ├── columns: a:3!null b:4!null
+      ├── left columns: a:11 b:12
+      ├── right columns: a:19 b:20
+      ├── cardinality: [0 - 1]
+      ├── key: ()
+      ├── fd: ()-->(3,4)
+      ├── distribution: east
+      ├── scan abc_part@b_idx
+      │    ├── columns: a:11!null b:12!null
+      │    ├── constraint: /9/12: [/'east'/1 - /'east'/1]
+      │    ├── cardinality: [0 - 1]
+      │    ├── key: ()
+      │    └── fd: ()-->(11,12)
+      └── scan abc_part@b_idx
+           ├── columns: a:19!null b:20!null
+           ├── constraint: /17/20
+           │    ├── [/'central'/1 - /'central'/1]
+           │    └── [/'west'/1 - /'west'/1]
+           ├── cardinality: [0 - 1]
+           ├── key: ()
+           └── fd: ()-->(19,20)
+
+# With optimizer_span_limit=3, the 3 spans are within the limit.
+opt locality=(region=east) expect=GenerateLocalityOptimizedScan set=(optimizer_span_limit=3)
+SELECT a FROM abc_part WHERE b = 1
+----
+project
+ ├── columns: a:3!null
+ ├── cardinality: [0 - 1]
+ ├── key: ()
+ ├── fd: ()-->(3)
+ ├── distribution: east
+ └── locality-optimized-search
+      ├── columns: a:3!null b:4!null
+      ├── left columns: a:11 b:12
+      ├── right columns: a:19 b:20
+      ├── cardinality: [0 - 1]
+      ├── key: ()
+      ├── fd: ()-->(3,4)
+      ├── distribution: east
+      ├── scan abc_part@b_idx
+      │    ├── columns: a:11!null b:12!null
+      │    ├── constraint: /9/12: [/'east'/1 - /'east'/1]
+      │    ├── cardinality: [0 - 1]
+      │    ├── key: ()
+      │    └── fd: ()-->(11,12)
+      └── scan abc_part@b_idx
+           ├── columns: a:19!null b:20!null
+           ├── constraint: /17/20
+           │    ├── [/'central'/1 - /'central'/1]
+           │    └── [/'west'/1 - /'west'/1]
+           ├── cardinality: [0 - 1]
+           ├── key: ()
+           └── fd: ()-->(19,20)
+
+# With optimizer_span_limit=2, the 3 spans exceed the limit, so the locality-
+# optimized scan does not fire.
+opt locality=(region=east) expect-not=GenerateLocalityOptimizedScan set=(optimizer_span_limit=2)
+SELECT a FROM abc_part WHERE b = 1
+----
+project
+ ├── columns: a:3!null
+ ├── cardinality: [0 - 1]
+ ├── key: ()
+ ├── fd: ()-->(3)
+ ├── distribution: east
+ └── select
+      ├── columns: a:3!null b:4!null
+      ├── cardinality: [0 - 1]
+      ├── key: ()
+      ├── fd: ()-->(3,4)
+      ├── distribution: east
+      ├── scan abc_part
+      │    ├── columns: a:3!null b:4
+      │    ├── check constraint expressions
+      │    │    ├── r:1 IN ('central', 'east', 'west') [outer=(1), constraints=(/1: [/'central' - /'central'] [/'east' - /'east'] [/'west' - /'west']; tight)]
+      │    │    └── t:2 IN (1, 2, 3) [outer=(2), constraints=(/2: [/1 - /1] [/2 - /2] [/3 - /3]; tight)]
+      │    ├── key: (3)
+      │    ├── fd: (3)-->(4), (4)~~>(3)
+      │    └── distribution: east
+      └── filters
+           └── b:4 = 1 [outer=(4), constraints=(/4: [/1 - /1]; tight), fd=()-->(4)]