opt: more accurate join multiplicity with FKs and self-joins

When calculating join multiplicity in the presence of a foreign key, it
was previously possible for the incorrect base table to be selected from
the right expression, in the case that the right expression contained a
self-join. The base table selected was incorrect because it did not
contain the equality column from the join filter. This prevented the
optimizer from determining that these joins preserve all rows from the
left side of the join, which could prevent further optimization of a
query.

Now, only the right equality columns are passed to `checkForeignKeyCase`
which ensures that the correct base table is selected. This is safe
because `verifyFiltersAreValidEqualities` has already checked that the
right equality columns are unfiltered in the right expression, so the
same checks in `checkForeignKeyCase` are redundant.

Release note (performance improvement): The optimizer better optimizes
queries that include both foreign key joins and self-joins.

pkg/sql/opt/memo/multiplicity_builder.go

@@ -122,8 +122,8 @@ func deriveUnfilteredCols(in RelExpr) opt.ColSet {
 		right := t.Child(1).(RelExpr)
 		multiplicity := GetJoinMultiplicity(t)
 
-		// Use the UnfilteredCols to determine whether unfiltered columns can be
-		// passed through.
+		// Use the join's multiplicity to determine whether unfiltered columns
+		// can be passed through.
 		if multiplicity.JoinPreservesLeftRows(t.Op()) {
 			unfilteredCols.UnionWith(deriveUnfilteredCols(left))
 		}
@@ -267,7 +267,8 @@ func filtersMatchAllLeftRows(left, right RelExpr, filters FiltersExpr) bool {
 			filters,
 		)
 	}
-	if !verifyFiltersAreValidEqualities(left, right, filters) {
+	rightEqualityCols, ok := verifyFiltersAreValidEqualities(left, right, filters)
+	if !ok {
 		return false
 	}
 	if checkSelfJoinCase(left.Memo().Metadata(), filters) {
@@ -276,41 +277,50 @@ func filtersMatchAllLeftRows(left, right RelExpr, filters FiltersExpr) bool {
 	}
 	// Case 2b.
 	return checkForeignKeyCase(
-		left.Memo().Metadata(), left.Relational().NotNullCols, deriveUnfilteredCols(right), filters)
+		left.Memo().Metadata(),
+		left.Relational().NotNullCols,
+		rightEqualityCols,
+		filters,
+	)
 }
 
-// verifyFiltersAreValidEqualities returns true when all of the following
-// conditions are satisfied:
-// 1. All filters are equalities.
-// 2. All equalities directly compare two columns.
-// 3. All equalities contain one column from the left not-null columns, and one
-//    column from the right unfiltered columns.
-// 4. All equality columns come from a base table.
-// 5. All left columns come from a single table, and all right columns come from
-//    a single table.
-func verifyFiltersAreValidEqualities(left, right RelExpr, filters FiltersExpr) bool {
+// verifyFiltersAreValidEqualities returns the set of equality columns in the
+// right relation and true when all of the following conditions are satisfied:
+//
+//   1. All filters are equalities.
+//   2. All equalities directly compare two columns.
+//   3. All equalities contain one column from the left not-null columns, and
+//      one column from the right unfiltered columns.
+//   4. All equality columns come from a base table.
+//   5. All left columns come from a single table, and all right columns come
+//      from a single table.
+//
+// Returns ok=false if any of these conditions are unsatisfied.
+func verifyFiltersAreValidEqualities(
+	left, right RelExpr, filters FiltersExpr,
+) (rightEqualityCols opt.ColSet, ok bool) {
 	md := left.Memo().Metadata()
 
 	var leftTab, rightTab opt.TableID
 	leftNotNullCols := left.Relational().NotNullCols
 	rightUnfilteredCols := deriveUnfilteredCols(right)
 	if rightUnfilteredCols.Empty() {
 		// There are no unfiltered columns from the right input.
-		return false
+		return opt.ColSet{}, false
 	}
 
 	for i := range filters {
 		eq, _ := filters[i].Condition.(*EqExpr)
 		if eq == nil {
 			// Condition #1: Conjunct is not an equality comparison.
-			return false
+			return opt.ColSet{}, false
 		}
 
 		leftVar, _ := eq.Left.(*VariableExpr)
 		rightVar, _ := eq.Right.(*VariableExpr)
 		if leftVar == nil || rightVar == nil {
 			// Condition #2: Conjunct does not directly compare two columns.
-			return false
+			return opt.ColSet{}, false
 		}
 
 		leftColID := leftVar.Col
@@ -322,7 +332,7 @@ func verifyFiltersAreValidEqualities(left, right RelExpr, filters FiltersExpr) b
 		}
 		if !leftNotNullCols.Contains(leftColID) || !rightUnfilteredCols.Contains(rightColID) {
 			// Condition #3: Columns don't come from both the left and right ColSets.
-			return false
+			return opt.ColSet{}, false
 		}
 
 		if leftTab == 0 || rightTab == 0 {
@@ -331,16 +341,19 @@ func verifyFiltersAreValidEqualities(left, right RelExpr, filters FiltersExpr) b
 			rightTab = md.ColumnMeta(rightColID).Table
 			if leftTab == 0 || rightTab == 0 {
 				// Condition #4: Columns don't come from base tables.
-				return false
+				return opt.ColSet{}, false
 			}
 		}
 		if leftTab != md.ColumnMeta(leftColID).Table || rightTab != md.ColumnMeta(rightColID).Table {
 			// Condition #5: The filter conditions reference more than one table from
 			// each side.
-			return false
+			return opt.ColSet{}, false
 		}
+
+		rightEqualityCols.Add(rightColID)
 	}
-	return true
+
+	return rightEqualityCols, true
 }
 
 // checkSelfJoinCase returns true if all equalities in the given FiltersExpr
@@ -374,7 +387,9 @@ func checkForeignKeyCase(
 ) bool {
 	if rightUnfilteredCols.Empty() {
 		// There are no unfiltered columns from the right; a valid foreign key
-		// relation is not possible.
+		// relation is not possible. This check, which is a duplicate of a check
+		// in verifyFiltersAreValidEqualities, is necessary in the case of a
+		// cross-join because verifyFiltersAreValidEqualities is not called.
 		return false
 	}
 
@@ -419,7 +434,7 @@ func checkForeignKeyCase(
 				rightColID := rightTableID.ColumnID(rightColOrd)
 				if !leftNotNullCols.Contains(leftColID) {
 					// The left column isn't a left not-null output column. It can't be
-					// used in a filter (since it isn't an output column) but the foreign key may still be valid.but it may still
+					// used in a filter (since it isn't an output column) but it may still
 					// be provably not null.
 					//
 					// A column that isn't in leftNotNullCols is guaranteed not-null if it
@@ -444,8 +459,11 @@ func checkForeignKeyCase(
 					continue
 				}
 				if !rightUnfilteredCols.Contains(rightColID) {
-					// The right column isn't guaranteed to be unfiltered. It can't be
-					// used in a filter.
+					// The right column isn't guaranteed to be unfiltered. It
+					// can't be used in a filter. This check, which is a
+					// duplicate of a check in verifyFiltersAreValidEqualities,
+					// is necessary in the case of a cross-join because
+					// verifyFiltersAreValidEqualities is not called.
 					continue
 				}
 				if filtersHaveEquality(filters, leftColID, rightColID) {

pkg/sql/opt/memo/multiplicity_builder_test.go

@@ -270,6 +270,36 @@ func TestGetJoinMultiplicity(t *testing.T) {
 			on:       ob.makeFilters(ob.makeEquality(fkCols[0], xyCols[0])),
 			expected: "left-rows(exactly-one)",
 		},
+		{ // 19
+			// SELECT *
+			// FROM fk_tab
+			// INNER JOIN (SELECT * FROM xy AS xy1 INNER JOIN xy xy2 ON xy1.x = xy2.x)
+			// ON r1 = xy1.x;
+			joinOp: opt.InnerJoinOp,
+			left:   fkScan,
+			right: ob.makeInnerJoin(
+				xyScan,
+				xyScan2,
+				ob.makeFilters(ob.makeEquality(xyCols[0], xyCols2[0])),
+			),
+			on:       ob.makeFilters(ob.makeEquality(fkCols[0], xyCols[0])),
+			expected: "left-rows(exactly-one), right-rows(zero-or-more)",
+		},
+		{ // 20
+			// SELECT *
+			// FROM fk_tab
+			// INNER JOIN (SELECT * FROM xy AS xy1 INNER JOIN xy xy2 ON xy1.x = xy2.x)
+			// ON r1 = xy2.x;
+			joinOp: opt.InnerJoinOp,
+			left:   fkScan,
+			right: ob.makeInnerJoin(
+				xyScan,
+				xyScan2,
+				ob.makeFilters(ob.makeEquality(xyCols[0], xyCols2[0])),
+			),
+			on:       ob.makeFilters(ob.makeEquality(fkCols[0], xyCols2[0])),
+			expected: "left-rows(exactly-one), right-rows(zero-or-more)",
+		},
 	}
 
 	for i, tc := range testCases {

pkg/sql/opt/memo/testdata/logprops/join

@@ -2224,7 +2224,7 @@ INNER JOIN (SELECT xysd.x, a.x AS t FROM xysd INNER JOIN xysd AS a ON xysd.x = a
 ----
 inner-join (hash)
  ├── columns: k:1(int!null) v:2(int) r1:3(int!null) r2:4(int) x:7(int!null) t:13(int!null)
- ├── multiplicity: left-rows(zero-or-one), right-rows(zero-or-more)
+ ├── multiplicity: left-rows(exactly-one), right-rows(zero-or-more)
  ├── key: (1)
  ├── fd: (1)-->(2-4), (7)==(3,13), (13)==(3,7), (3)==(7,13)
  ├── prune: (1,2,4)