opt: respect NO_INDEX_JOIN flag

Prior to this commit, it was possible that the optimizer could
produce a plan with an index join even if the user hinted that
index joins should be avoided by using the NO_INDEX_JOIN hint. This
commit fixes that oversight, and we no longer plan an index join
in this case. This commit also adds assertions that an index join
is not planned if NO_INDEX_JOIN is used to prevent this bug from
recurring.

Fixes #85841

Release note (bug fix): Fixed an issue where the NO_INDEX_JOIN
hint could be ignored by the optimizer in some cases, causing it
to create a query plan with an index join.

pkg/sql/opt/memo/check_expr.go

@@ -207,6 +207,9 @@ func (m *Memo) CheckExpr(e opt.Expr) {
 		if t.Cols.Empty() {
 			panic(errors.AssertionFailedf("index join with no columns"))
 		}
+		if scan, ok := t.Input.(*ScanExpr); ok && scan.Flags.NoIndexJoin {
+			panic(errors.AssertionFailedf("index join used with NoIndexJoin flag"))
+		}
 
 	case *LookupJoinExpr:
 		if len(t.KeyCols) == 0 && len(t.LookupExpr) == 0 {

pkg/sql/opt/xform/groupby_funcs.go

@@ -411,6 +411,12 @@ func (c *CustomFuncs) GenerateLimitedGroupByScans(
 			return
 		}
 
+		// Otherwise, try to construct an IndexJoin operator that provides the
+		// columns missing from the index.
+		if sp.Flags.NoIndexJoin {
+			return
+		}
+
 		// Calculate the PK columns once.
 		if pkCols.Empty() {
 			pkCols = c.PrimaryKeyCols(sp.Table)

pkg/sql/opt/xform/index_scan_builder.go

@@ -163,6 +163,9 @@ func (b *indexScanBuilder) AddSelectAfterSplit(
 // AddIndexJoin wraps the input expression with an IndexJoin expression that
 // produces the given set of columns by lookup in the primary index.
 func (b *indexScanBuilder) AddIndexJoin(cols opt.ColSet) {
+	if b.scanPrivate.Flags.NoIndexJoin {
+		panic(errors.AssertionFailedf("attempt to create an index join with NoIndexJoin flag"))
+	}
 	if b.hasIndexJoin() {
 		panic(errors.AssertionFailedf("cannot call AddIndexJoin twice"))
 	}

pkg/sql/opt/xform/limit_funcs.go

@@ -68,7 +68,7 @@ func (c *CustomFuncs) CanLimitFilteredScan(
 	if scanPrivate.IsVirtualTable(md) && !required.Any() {
 		return false
 	}
-	ok, _ := ordering.ScanPrivateCanProvide(c.e.mem.Metadata(), scanPrivate, &required)
+	ok, _ := ordering.ScanPrivateCanProvide(md, scanPrivate, &required)
 	return ok
 }
 
@@ -298,6 +298,12 @@ func (c *CustomFuncs) GenerateLimitedTopKScans(
 			return
 		}
 
+		// Otherwise, try to construct an IndexJoin operator that provides the
+		// columns missing from the index.
+		if sp.Flags.NoIndexJoin {
+			return
+		}
+
 		// Calculate the PK columns once.
 		if pkCols.Empty() {
 			pkCols = c.PrimaryKeyCols(sp.Table)

pkg/sql/opt/xform/select_funcs.go

@@ -119,6 +119,12 @@ func (c *CustomFuncs) GeneratePartialIndexScans(
 			return
 		}
 
+		// Otherwise, try to construct an IndexJoin operator that provides the
+		// columns missing from the index.
+		if scanPrivate.Flags.NoIndexJoin {
+			return
+		}
+
 		// Calculate the PK columns once.
 		if pkCols.Empty() {
 			pkCols = c.PrimaryKeyCols(scanPrivate.Table)
@@ -853,6 +859,10 @@ func (c *CustomFuncs) GenerateInvertedIndexScans(
 		newScanPrivate.SetConstraint(c.e.evalCtx, constraint)
 		newScanPrivate.InvertedConstraint = spansToRead
 
+		if scanPrivate.Flags.NoIndexJoin {
+			return
+		}
+
 		// Calculate the PK columns once.
 		if pkCols.Empty() {
 			pkCols = c.PrimaryKeyCols(scanPrivate.Table)

pkg/sql/opt/xform/testdata/rules/groupby

@@ -3517,3 +3517,28 @@ limit
  │    └── aggregations
  │         └── count-rows [as=count_rows:8]
  └── 10
+
+# GenerateLimitedGroupByScans will be triggered, but not add an index
+# scan to the memo since NO_INDEX_JOIN is specified.
+memo expect-not=GenerateLimitedGroupByScans
+SELECT d, e, count(*) FROM defg@{NO_INDEX_JOIN} GROUP BY d, e LIMIT 10
+----
+memo (optimized, ~5KB, required=[presentation: d:1,e:2,count:8])
+ ├── G1: (limit G2 G3)
+ │    └── [presentation: d:1,e:2,count:8]
+ │         ├── best: (limit G2="[limit hint: 10.00]" G3)
+ │         └── cost: 1145.01
+ ├── G2: (group-by G4 G5 cols=(1,2)) (group-by G4 G5 cols=(1,2),ordering=+1)
+ │    └── [limit hint: 10.00]
+ │         ├── best: (group-by G4 G5 cols=(1,2))
+ │         └── cost: 1144.90
+ ├── G3: (const 10)
+ ├── G4: (scan defg,cols=(1,2))
+ │    ├── [ordering: +1] [limit hint: 10.00]
+ │    │    ├── best: (sort G4)
+ │    │    └── cost: 1334.20
+ │    └── []
+ │         ├── best: (scan defg,cols=(1,2))
+ │         └── cost: 1094.72
+ ├── G5: (aggregations G6)
+ └── G6: (count-rows)

pkg/sql/opt/xform/testdata/rules/limit

@@ -1936,6 +1936,34 @@ top-k
  └── scan defg
       └── columns: d:1 e:2 f:3 g:4
 
+# GenerateLimitedTopKScans will be triggered, but not add an index
+# scan to the memo since NO_INDEX_JOIN is specified.
+memo expect-not=GenerateLimitedTopKScans
+SELECT d, f, e FROM defg@{NO_INDEX_JOIN} ORDER BY d, f, e LIMIT 10
+----
+memo (optimized, ~4KB, required=[presentation: d:1,f:3,e:2] [ordering: +1,+3,+2])
+ ├── G1: (limit G2 G3 ordering=+1,+3,+2) (top-k G2 &{10 +1,+3,+2 }) (top-k G2 &{10 +1,+3,+2 +1,+3})
+ │    ├── [presentation: d:1,f:3,e:2] [ordering: +1,+3,+2]
+ │    │    ├── best: (top-k G2 &{10 +1,+3,+2 })
+ │    │    └── cost: 1196.32
+ │    └── []
+ │         ├── best: (top-k G2 &{10 +1,+3,+2 })
+ │         └── cost: 1196.32
+ ├── G2: (scan defg,cols=(1-3))
+ │    ├── [ordering: +1,+3,+2] [limit hint: 10.00]
+ │    │    ├── best: (sort G2)
+ │    │    └── cost: 1366.36
+ │    ├── [ordering: +1,+3]
+ │    │    ├── best: (sort G2)
+ │    │    └── cost: 1365.27
+ │    ├── [ordering: +1,+3] [limit hint: 100.00]
+ │    │    ├── best: (sort G2)
+ │    │    └── cost: 1365.27
+ │    └── []
+ │         ├── best: (scan defg,cols=(1-3))
+ │         └── cost: 1104.82
+ └── G3: (const 10)
+
 # ---------------------------------------------------
 # GeneratePartialOrderTopK
 # ---------------------------------------------------

pkg/sql/opt/xform/testdata/rules/select

@@ -306,6 +306,29 @@ memo (optimized, ~17KB, required=[presentation: k:1,i:2,f:3,s:4,b:5])
  ├── G16: (variable s)
  └── G17: (const 'foo')
 
+# GeneratePartialIndexScans will be triggered, but not add an index
+# scan to the memo since NO_INDEX_JOIN is specified.
+memo expect-not=GeneratePartialIndexScans
+SELECT * FROM p@{NO_INDEX_JOIN} WHERE i > 0 AND s = 'foo'
+----
+memo (optimized, ~9KB, required=[presentation: k:1,i:2,f:3,s:4,b:5])
+ ├── G1: (select G2 G3)
+ │    └── [presentation: k:1,i:2,f:3,s:4,b:5]
+ │         ├── best: (select G2 G3)
+ │         └── cost: 1135.06
+ ├── G2: (scan p,cols=(1-5))
+ │    └── []
+ │         ├── best: (scan p,cols=(1-5))
+ │         └── cost: 1125.02
+ ├── G3: (filters G4 G5)
+ ├── G4: (gt G6 G7)
+ ├── G5: (eq G8 G9)
+ ├── G6: (variable i)
+ ├── G7: (const 0)
+ ├── G8: (variable s)
+ └── G9: (const 'foo')
+
+
 # Do not generate a partial index scan when the predicate is not implied by the
 # filter.
 memo expect-not=GeneratePartialIndexScans
@@ -2200,6 +2223,30 @@ memo (optimized, ~8KB, required=[presentation: k:1])
  ├── G8: (variable j)
  └── G9: (const '{"a": "b"}')
 
+# GenerateInvertedIndexScans will be triggered, but not add an index
+# scan to the memo since NO_INDEX_JOIN is specified.
+memo expect-not=GenerateInvertedIndexScans
+SELECT k FROM b@{NO_INDEX_JOIN} WHERE j @> '{"a": "b"}'
+----
+memo (optimized, ~6KB, required=[presentation: k:1])
+ ├── G1: (project G2 G3 k)
+ │    └── [presentation: k:1]
+ │         ├── best: (project G2 G3 k)
+ │         └── cost: 1095.77
+ ├── G2: (select G4 G5)
+ │    └── []
+ │         ├── best: (select G4 G5)
+ │         └── cost: 1094.65
+ ├── G3: (projections)
+ ├── G4: (scan b,cols=(1,4))
+ │    └── []
+ │         ├── best: (scan b,cols=(1,4))
+ │         └── cost: 1084.62
+ ├── G5: (filters G6)
+ ├── G6: (contains G7 G8)
+ ├── G7: (variable j)
+ └── G8: (const '{"a": "b"}')
+
 # Query requiring an index join with no remaining filter.
 opt expect=GenerateInvertedIndexScans
 SELECT u, k FROM b WHERE j @> '{"a": "b"}'