xform: fix missing derived terms from cockroachdb#90599

This commit fixes an oversight in cockroachdb#90599, which derives t1.crdb_region = t2.crdb_region
join terms from FK constraints for lookup join. That PR only examined strict join keys
in the lookup relation. However, the parent table may be the input relation. The fix
is to also perform the strict join key check with the input and lookup tables swapped
when building a lookup join.

Release note: None

pkg/ccl/logictestccl/testdata/logic_test/regional_by_row_query_behavior

@@ -3021,35 +3021,33 @@ CREATE TABLE user_settings_cascades (
 # user_settings and avoid visiting all regions.
 query T retry
 EXPLAIN SELECT users.crdb_region AS user_region, user_settings.crdb_region AS user_settings_region, *
-  FROM users JOIN user_settings ON users.id = user_settings.user_id AND users.id = '5ebfedee-0dcf-41e6-a315-5fa0b51b9882';
+  FROM users INNER LOOKUP JOIN user_settings ON users.id = user_settings.user_id AND users.id = '5ebfedee-0dcf-41e6-a315-5fa0b51b9882';
 ----
 distribution: local
 vectorized: true
 ·
-• merge join
-│ equality: (user_id) = (id)
-│ right cols are key
-│
-├── • index join
-│   │ table: user_settings@user_settings_pkey
-│   │
-│   └── • scan
-│         missing stats
-│         table: user_settings@user_settings_user_id_idx
-│         spans: [/'ap-southeast-2'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'ap-southeast-2'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882'] [/'ca-central-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'ca-central-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882'] [/'us-east-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'us-east-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882']
+• lookup join
+│ table: user_settings@user_settings_pkey
+│ equality: (crdb_region, id) = (crdb_region,id)
+│ equality cols are key
 │
-└── • union all
-    │ limit: 1
+└── • lookup join
+    │ table: user_settings@user_settings_user_id_idx
+    │ equality: (crdb_region, id) = (crdb_region,user_id)
+    │ pred: user_id = '5ebfedee-0dcf-41e6-a315-5fa0b51b9882'
     │
-    ├── • scan
-    │     missing stats
-    │     table: users@users_pkey
-    │     spans: [/'ap-southeast-2'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'ap-southeast-2'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882']
-    │
-    └── • scan
-          missing stats
-          table: users@users_pkey
-          spans: [/'ca-central-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'ca-central-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882'] [/'us-east-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'us-east-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882']
+    └── • union all
+        │ limit: 1
+        │
+        ├── • scan
+        │     missing stats
+        │     table: users@users_pkey
+        │     spans: [/'ap-southeast-2'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'ap-southeast-2'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882']
+        │
+        └── • scan
+              missing stats
+              table: users@users_pkey
+              spans: [/'ca-central-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'ca-central-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882'] [/'us-east-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882' - /'us-east-1'/'5ebfedee-0dcf-41e6-a315-5fa0b51b9882']
 
 # Ensure that the FK checks and cascades are efficient.
 query T

pkg/sql/opt/distribution/distribution.go

@@ -158,7 +158,7 @@ func BuildLookupJoinLookupTableDistribution(
 	lookupTableMeta := lookupJoin.Memo().Metadata().TableMeta(lookupJoin.Table)
 	lookupTable := lookupTableMeta.Table
 
-	if lookupJoin.LocalityOptimized {
+	if lookupJoin.LocalityOptimized || lookupJoin.ChildOfLocalityOptimizedSearch {
 		provided.FromLocality(evalCtx.Locality)
 		return provided
 	} else if lookupTable.IsGlobalTable() {
@@ -190,6 +190,10 @@ func BuildLookupJoinLookupTableDistribution(
 			if filterIdx, ok := lookupJoin.GetConstPrefixFilter(lookupJoin.Memo().Metadata()); ok {
 				firstIndexColEqExpr := lookupJoin.LookupJoinPrivate.LookupExpr[filterIdx].Condition
 				if firstIndexColEqExpr.Op() == opt.EqOp {
+					// TODO(msirek): For a join between 2 REGIONAL BY ROW tables, check if
+					//               there is a crdb_region = crdb_region term. If there
+					//               is, use the Distribution of the input table as the
+					//               Distribution of the lookup table.
 					if regionName, ok := GetDEnumAsStringFromConstantExpr(firstIndexColEqExpr.Child(1)); ok {
 						provided.Regions = []string{regionName}
 						return provided

pkg/sql/opt/xform/join_funcs.go

@@ -385,12 +385,36 @@ func (c *CustomFuncs) generateLookupJoinsImpl(
 	onClauseLookupRelStrictKeyCols, lookupRelEquijoinCols, inputRelJoinCols, lookupIsKey :=
 		c.GetEquijoinStrictKeyCols(on, scanPrivate, input)
 
+	var input2 memo.RelExpr
+	var scanPrivate2 *memo.ScanPrivate
+	var lookupIsKey2 bool
+	var indexCols2 opt.ColSet
+	// Look up strict key cols in the reverse direction for the purposes of
+	// deriving a `t1.crdb_region = t2.crdb_region` term based on foreign key
+	// constraints.
+	if !lookupIsKey {
+		possibleScanExpr := input
+		if selectExpr, ok := input.(*memo.SelectExpr); ok {
+			possibleScanExpr = selectExpr.Input
+		}
+		if scanExpr, ok := possibleScanExpr.(*memo.ScanExpr); ok {
+			scanPrivate2 = &scanExpr.ScanPrivate
+			// The scan should already exist in the memo. We need to look it up so we
+			// have a `ScanExpr` with properties fully populated.
+			input2 = scanExpr.Memo().MemoizeScan(scanPrivate)
+			tabMeta := c.e.mem.Metadata().TableMeta(scanPrivate2.Table)
+			indexCols2 = tabMeta.IndexColumns(scanPrivate2.Index)
+			onClauseLookupRelStrictKeyCols, lookupRelEquijoinCols, inputRelJoinCols, lookupIsKey2 =
+				c.GetEquijoinStrictKeyCols(on, scanPrivate2, input2)
+		}
+	}
+
 	var pkCols opt.ColList
 	var newScanPrivate *memo.ScanPrivate
 	var iter scanIndexIter
 	iter.Init(c.e.evalCtx, c.e.f, c.e.mem, &c.im, scanPrivate, on, rejectInvertedIndexes)
 	iter.ForEach(func(index cat.Index, onFilters memo.FiltersExpr, indexCols opt.ColSet, _ bool, _ memo.ProjectionsExpr) {
-		// Skip indexes that do no cover all virtual projection columns, if
+		// Skip indexes that do not cover all virtual projection columns, if
 		// there are any. This can happen when there are multiple virtual
 		// columns indexed in different indexes.
 		//
@@ -405,6 +429,9 @@ func (c *CustomFuncs) generateLookupJoinsImpl(
 		if lookupIsKey {
 			derivedfkOnFilters = c.ForeignKeyConstraintFilters(
 				input, scanPrivate, indexCols, onClauseLookupRelStrictKeyCols, lookupRelEquijoinCols, inputRelJoinCols)
+		} else if lookupIsKey2 {
+			derivedfkOnFilters = c.ForeignKeyConstraintFilters(
+				input2, scanPrivate2, indexCols2, onClauseLookupRelStrictKeyCols, lookupRelEquijoinCols, inputRelJoinCols)
 		}
 		lookupConstraint, foundEqualityCols := cb.Build(index, onFilters, optionalFilters, derivedfkOnFilters)
 		if lookupConstraint.IsUnconstrained() {