opt: improve EnsureKey custom function

This commit makes a small improvement to the `EnsureKey` custom function
(used in decorrelation rules), so that it can add passthrough columns to
a `Project` operator in an effort to find a key. This can prevent rules
from adding unnecessary `Ordinality` operators to the query plan.

Epic: None

Release note: None

pkg/sql/opt/memo/testdata/typing

@@ -419,18 +419,16 @@ SELECT EXISTS(SELECT * FROM a WHERE expr<0) FROM (SELECT x+1 AS expr FROM a)
 project
  ├── columns: exists:11(bool!null)
  ├── group-by (hash)
- │    ├── columns: true_agg:13(bool) rownum:15(int!null)
- │    ├── grouping columns: rownum:15(int!null)
+ │    ├── columns: x:1(int!null) true_agg:13(bool)
+ │    ├── grouping columns: x:1(int!null)
  │    ├── left-join (cross)
- │    │    ├── columns: expr:5(int!null) true:12(bool) rownum:15(int!null)
- │    │    ├── ordinality
- │    │    │    ├── columns: expr:5(int!null) rownum:15(int!null)
- │    │    │    └── project
- │    │    │         ├── columns: expr:5(int!null)
- │    │    │         ├── scan a
- │    │    │         │    └── columns: x:1(int!null)
- │    │    │         └── projections
- │    │    │              └── x:1 + 1 [as=expr:5, type=int]
+ │    │    ├── columns: x:1(int!null) expr:5(int!null) true:12(bool)
+ │    │    ├── project
+ │    │    │    ├── columns: expr:5(int!null) x:1(int!null)
+ │    │    │    ├── scan a
+ │    │    │    │    └── columns: x:1(int!null)
+ │    │    │    └── projections
+ │    │    │         └── x:1 + 1 [as=expr:5, type=int]
  │    │    ├── project
  │    │    │    ├── columns: true:12(bool!null)
  │    │    │    ├── scan a

pkg/sql/opt/norm/decorrelate_funcs.go

@@ -527,14 +527,9 @@ func (c *CustomFuncs) ConstructApplyJoin(
 // input expression (perhaps augmented with a key column(s) or wrapped by
 // Ordinality).
 func (c *CustomFuncs) EnsureKey(in memo.RelExpr) memo.RelExpr {
-	_, ok := c.CandidateKey(in)
-	if ok {
-		return in
-	}
-
 	// Try to add the preexisting primary key if the input is a Scan or Scan
 	// wrapped in a Select.
-	if res, ok := c.TryAddKeyToScan(in); ok {
+	if res, ok := c.tryFindExistingKey(in); ok {
 		return res
 	}
 
@@ -544,34 +539,34 @@ func (c *CustomFuncs) EnsureKey(in memo.RelExpr) memo.RelExpr {
 	return c.f.ConstructOrdinality(in, &private)
 }
 
-// TryAddKeyToScan checks whether the input expression is a non-virtual table
-// Scan, either alone or wrapped in a Select. If so, it returns a new Scan
-// (possibly wrapped in a Select) augmented with the preexisting primary key
-// for the table.
-func (c *CustomFuncs) TryAddKeyToScan(in memo.RelExpr) (_ memo.RelExpr, ok bool) {
-	augmentScan := func(scan *memo.ScanExpr) (_ memo.RelExpr, ok bool) {
-		private := scan.ScanPrivate
+// tryFindExistingKey attempts to find an existing key for the input expression.
+// It may modify the expression in order to project the key column.
+func (c *CustomFuncs) tryFindExistingKey(in memo.RelExpr) (_ memo.RelExpr, ok bool) {
+	_, hasKey := c.CandidateKey(in)
+	if hasKey {
+		return in, true
+	}
+	switch t := in.(type) {
+	case *memo.ProjectExpr:
+		input, foundKey := c.tryFindExistingKey(t.Input)
+		if foundKey {
+			return c.f.ConstructProject(input, t.Projections, input.Relational().OutputCols), true
+		}
+
+	case *memo.ScanExpr:
+		private := t.ScanPrivate
 		tableID := private.Table
 		table := c.f.Metadata().Table(tableID)
 		if !table.IsVirtualTable() {
 			keyCols := c.PrimaryKeyCols(tableID)
 			private.Cols = private.Cols.Union(keyCols)
 			return c.f.ConstructScan(&private), true
 		}
-		return nil, false
-	}
-
-	switch t := in.(type) {
-	case *memo.ScanExpr:
-		if res, ok := augmentScan(t); ok {
-			return res, true
-		}
 
 	case *memo.SelectExpr:
-		if scan, ok := t.Input.(*memo.ScanExpr); ok {
-			if res, ok := augmentScan(scan); ok {
-				return c.f.ConstructSelect(res, t.Filters), true
-			}
+		input, foundKey := c.tryFindExistingKey(t.Input)
+		if foundKey {
+			return c.f.ConstructSelect(input, t.Filters), true
 		}
 	}
 

pkg/sql/opt/norm/testdata/rules/decorrelate

@@ -5417,26 +5417,25 @@ project
  ├── columns: r:13
  ├── immutable
  ├── group-by (hash)
- │    ├── columns: scalar:14!null bool_or:16 rownum:18!null
- │    ├── grouping columns: rownum:18!null
+ │    ├── columns: k:1!null scalar:14!null bool_or:16
+ │    ├── grouping columns: k:1!null
  │    ├── immutable
- │    ├── key: (18)
- │    ├── fd: ()-->(14), (18)-->(14,16)
+ │    ├── key: (1)
+ │    ├── fd: ()-->(14), (1)-->(14,16)
  │    ├── left-join-apply
- │    │    ├── columns: i:2 column12:12 scalar:14!null notnull:15 rownum:18!null
+ │    │    ├── columns: k:1!null i:2 column12:12 scalar:14!null notnull:15
  │    │    ├── immutable
- │    │    ├── fd: ()-->(14), (18)-->(2)
- │    │    ├── ordinality
- │    │    │    ├── columns: i:2 scalar:14!null rownum:18!null
- │    │    │    ├── key: (18)
- │    │    │    ├── fd: ()-->(14), (18)-->(2,14)
- │    │    │    └── project
- │    │    │         ├── columns: scalar:14!null i:2
- │    │    │         ├── fd: ()-->(14)
- │    │    │         ├── scan a
- │    │    │         │    └── columns: i:2
- │    │    │         └── projections
- │    │    │              └── (5, 50) [as=scalar:14]
+ │    │    ├── fd: ()-->(14), (1)-->(2)
+ │    │    ├── project
+ │    │    │    ├── columns: scalar:14!null k:1!null i:2
+ │    │    │    ├── key: (1)
+ │    │    │    ├── fd: ()-->(14), (1)-->(2)
+ │    │    │    ├── scan a
+ │    │    │    │    ├── columns: k:1!null i:2
+ │    │    │    │    ├── key: (1)
+ │    │    │    │    └── fd: (1)-->(2)
+ │    │    │    └── projections
+ │    │    │         └── (5, 50) [as=scalar:14]
  │    │    ├── project
  │    │    │    ├── columns: notnull:15!null column12:12!null
  │    │    │    ├── outer: (2)

pkg/sql/opt/norm/testdata/rules/inline

@@ -1042,28 +1042,25 @@ project
  ├── columns: exists:14!null
  ├── immutable
  ├── group-by (hash)
- │    ├── columns: true_agg:16 rownum:18!null
- │    ├── grouping columns: rownum:18!null
+ │    ├── columns: k:1!null true_agg:16
+ │    ├── grouping columns: k:1!null
  │    ├── immutable
- │    ├── key: (18)
- │    ├── fd: (18)-->(16)
+ │    ├── key: (1)
+ │    ├── fd: (1)-->(16)
  │    ├── left-join (cross)
- │    │    ├── columns: expr:8!null true:15 rownum:18!null
+ │    │    ├── columns: k:1!null expr:8!null true:15
  │    │    ├── immutable
- │    │    ├── fd: (18)-->(8)
- │    │    ├── ordinality
- │    │    │    ├── columns: expr:8!null rownum:18!null
+ │    │    ├── fd: (1)-->(8)
+ │    │    ├── project
+ │    │    │    ├── columns: expr:8!null k:1!null
  │    │    │    ├── immutable
- │    │    │    ├── key: (18)
- │    │    │    ├── fd: (18)-->(8)
- │    │    │    └── project
- │    │    │         ├── columns: expr:8!null
- │    │    │         ├── immutable
- │    │    │         ├── scan a
- │    │    │         │    ├── columns: k:1!null
- │    │    │         │    └── key: (1)
- │    │    │         └── projections
- │    │    │              └── k:1 + 1 [as=expr:8, outer=(1), immutable]
+ │    │    │    ├── key: (1)
+ │    │    │    ├── fd: (1)-->(8)
+ │    │    │    ├── scan a
+ │    │    │    │    ├── columns: k:1!null
+ │    │    │    │    └── key: (1)
+ │    │    │    └── projections
+ │    │    │         └── k:1 + 1 [as=expr:8, outer=(1), immutable]
  │    │    ├── project
  │    │    │    ├── columns: true:15!null
  │    │    │    ├── fd: ()-->(15)