Merge #147490

147490: opt: do not produce unneeded PK cols in inverted index scans r=mgartner a=mgartner

A minor bug has been fixed where an inverted index scan could produce
more PK columns than necessary. This caused test-only assertion failures
but does not cause any known, user-visible bugs. The bug was caused by
the incorrect assumption that all inverted index scans need to produces
all PK columns. In actuality, an inverted index scan can produce a
subset of PK columns if the original scan produces a subset of PK
columns and neither an inverted filter nor an index join are needed.

Fixes #143070

Release note: None


Co-authored-by: Marcus Gartner <[email protected]>

Author: craig[bot]

pkg/sql/opt/xform/select_funcs.go

@@ -968,17 +968,33 @@ func (c *CustomFuncs) generateInvertedIndexScansImpl(
 			pkCols = c.PrimaryKeyCols(scanPrivate.Table)
 		}
 
+		// Start with a new Scan that produces only the PK columns produced by
+		// the original scan.
+		// NOTE: Intersection is used intentionally to avoid mutating pkCols.
+		newScanPrivate.Cols = pkCols.Intersection(scanPrivate.Cols)
+
+		// We will need an index join above the scan if the original scan
+		// produces non-PK columns.
+		needIndexJoin := !scanPrivate.Cols.SubsetOf(newScanPrivate.Cols)
 		// We will need an inverted filter above the scan if the spanExpr might
 		// produce duplicate primary keys or requires at least one UNION or
 		// INTERSECTION. In this case, we must scan both the primary key columns
 		// and the inverted key column.
 		needInvertedFilter := !spanExpr.Unique || spanExpr.Operator != inverted.None
-		newScanPrivate.Cols = pkCols.Copy()
+
+		// An index join or an inverted filter require all PK columns.
+		if needIndexJoin || needInvertedFilter {
+			newScanPrivate.Cols.UnionWith(pkCols)
+		}
+
 		var invertedCol opt.ColumnID
 		if needInvertedFilter {
+			// If we need an inverted filter, then we must also produce the
+			// inverted key column.
 			invertedCol = scanPrivate.Table.ColumnID(index.InvertedColumn().Ordinal())
 			newScanPrivate.Cols.Add(invertedCol)
 		}
+
 		sb.SetScan(&newScanPrivate)
 
 		// Add an inverted filter if needed.
@@ -990,9 +1006,7 @@ func (c *CustomFuncs) generateInvertedIndexScansImpl(
 		// be applied above the scan, and one that requires columns not produced
 		// by the scan.
 		filters = sb.AddSelectAfterSplit(filters, pkCols)
-		if !scanPrivate.Cols.SubsetOf(newScanPrivate.Cols) {
-			// Add an index join if the scan does not produce all the needed
-			// columns.
+		if needIndexJoin {
 			sb.AddIndexJoin(scanPrivate.Cols)
 		}
 		// Add the remaining filters, if any.

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

@@ -8567,7 +8567,7 @@ CREATE TABLE t136078 (
 )
 ----
 
-opt set=(optimizer_use_trigram_similarity_optimization=false)
+opt expect=GenerateInvertedIndexScans set=(optimizer_use_trigram_similarity_optimization=false)
 SELECT b FROM t136078 WHERE 'abc' % col3_8
 ----
 project
@@ -8598,6 +8598,36 @@ project
            └── filters
                 └── 'abc' % lower(b:2) [outer=(2), stable]
 
+exec-ddl
+CREATE TABLE t143070 (
+  a STRING,
+  b STRING,
+  c STRING AS (b) VIRTUAL,
+  PRIMARY KEY (a, b),
+  INVERTED INDEX (c gin_trgm_ops)
+)
+----
+
+# Regression test for #143070. The inverted index scan should not produce
+# unneeded PK cols.
+opt expect=GenerateInvertedIndexScans
+SELECT 1 FROM t143070 WHERE c LIKE 'aaa'
+----
+project
+ ├── columns: "?column?":7!null
+ ├── fd: ()-->(7)
+ ├── select
+ │    ├── columns: b:2!null
+ │    ├── fd: ()-->(2)
+ │    ├── scan t143070@t143070_c_idx,inverted
+ │    │    ├── columns: b:2!null
+ │    │    └── inverted constraint: /6/1/2
+ │    │         └── spans: ["aaa", "aaa"]
+ │    └── filters
+ │         └── b:2 LIKE 'aaa' [outer=(2), constraints=(/2: [/'aaa' - /'aaa']; tight), fd=()-->(2)]
+ └── projections
+      └── 1 [as="?column?":7]
+
 
 # --------------------------------------------------
 # GenerateMinimalInvertedIndexScans