xform: compute distribution of a lookup join with a lookup join as input

The adjustment to optimizer costing of lookup joins, added in v23.1,
to account for non-local input relation distribution (among the regions
in a multiregion database) did not account for the case where the
input to the lookup join is itself a lookup join. This resulted in an
extra distribution cost being added to some best-cost lookup joins,
causing the optimizer to no longer select them.

The issue is in function `getCRBDRegionColFromInput`, which only looks
for scan and locality-optimized scan operations as inputs. It is now
updated with a recursive call to itself plus a call to
`BuildLookupJoinLookupTableDistribution`, so the proper distribution
of a chain of lookup joins, plus the associated `crdb_region` column
in the lookup table can be determined. Function
`BuildLookupJoinLookupTableDistribution` is updated to return the
first lookup index column, when it is matched with the input's
crdb_region column, to facilitate this behavior.

Fixes: cockroachdb#105942

Release note (bug fix): This patch fixes an optimizer costing bug
introduced in v23.1 which may cause a query involving 2 or more
joins with REGIONAL BY ROW tables to not pick the most optimal
lookup joins.

Author: Mark Sirek

pkg/ccl/logictestccl/testdata/logic_test/multi_region_remote_access_error

@@ -639,7 +639,7 @@ locality-optimized-search
 # Locality optimized search with lookup join will be supported in phase 2 or 3
 # when we can dynamically determine if the lookup will access a remote region.
 retry
-statement error pq: Query has no home region\. Try adding a filter on o\.crdb_region and/or on key column \(o\.id\)\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+statement error pq: Query has no home region\. Try adding a filter on o\.crdb_region and/or on key column \(o\.cust_id\)\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
 SELECT * FROM customers c JOIN orders o ON c.id = o.cust_id AND
   (c.crdb_region = o.crdb_region) WHERE c.id = '69a1c2c2-5b18-459e-94d2-079dc53a4dd0'
 

pkg/ccl/logictestccl/testdata/logic_test/regional_by_row_query_behavior

@@ -3554,3 +3554,91 @@ EXPLAIN (VEC) SELECT * FROM child LEFT JOIN parent ON p_id = c_p_id WHERE c_id =
 
 statement ok
 RESET vectorize
+
+statement ok
+CREATE TABLE abc (
+  id UUID PRIMARY KEY,
+  id1 UUID NOT NULL,
+  created_at TIMESTAMP NOT NULL,
+  updated_at TIMESTAMP NOT NULL,
+  id2 UUID NULL,
+  INDEX (id1 ASC, id2 ASC)
+) LOCALITY REGIONAL BY ROW;
+
+statement ok
+CREATE TABLE xyz (
+  id UUID PRIMARY KEY,
+  str STRING NOT NULL,
+  abc_id UUID NOT NULL,
+  id2 UUID NULL,
+  FOREIGN KEY (abc_id) REFERENCES abc(id) ON DELETE CASCADE,
+  INDEX (abc_id ASC, id2 ASC),
+  INDEX (id2 ASC, str ASC, abc_id ASC)
+) LOCALITY REGIONAL BY ROW;
+
+statement ok
+ALTER TABLE abc INJECT STATISTICS '[
+  {
+    "avg_size": 3000,
+    "columns": ["id"],
+    "distinct_count": 100,
+    "row_count": 100,
+    "created_at": "2018-01-01 1:00:00.00000+00:00"
+  }
+]';
+
+statement ok
+ALTER TABLE xyz INJECT STATISTICS '[
+  {
+    "avg_size": 3000,
+    "columns": ["id"],
+    "distinct_count": 100,
+    "row_count": 100,
+    "created_at": "2018-01-01 1:00:00.00000+00:00"
+  }
+]';
+
+# Regression test for #105942
+# The following should produce 2 lookup joins with a local distribution.
+query T retry
+EXPLAIN SELECT
+  xyz.str,
+  abc.id,
+  abc.id1,
+  abc.id2,
+  abc.created_at,
+  abc.updated_at
+FROM
+  abc JOIN xyz ON
+    xyz.abc_id = abc.id
+    AND xyz.id2 = abc.id2
+    AND xyz.crdb_region = abc.crdb_region
+WHERE
+  abc.id1 = '6da4f356-e526-4b78-b9f9-bbb1a7fc12d6'
+  AND abc.id2 = '68088706-02c6-47d1-b993-a421cd761f2b'
+  AND abc.crdb_region = 'ap-southeast-2'
+  AND xyz.crdb_region = 'ap-southeast-2';
+----
+distribution: local
+vectorized: true
+·
+• lookup join
+│ estimated row count: 1
+│ table: abc@abc_pkey
+│ equality: (crdb_region, id) = (crdb_region,id)
+│ equality cols are key
+│
+└── • lookup join
+    │ estimated row count: 0
+    │ table: abc@abc_id1_id2_idx
+    │ equality: (crdb_region, lookup_join_const_col_@2, id2, abc_id) = (crdb_region,id1,id2,id)
+    │ equality cols are key
+    │ pred: (id2 = '68088706-02c6-47d1-b993-a421cd761f2b') AND (crdb_region = 'ap-southeast-2')
+    │
+    └── • render
+        │
+        └── • scan
+              estimated row count: 4 (3.7% of the table; stats collected <hidden> ago)
+              table: xyz@xyz_id2_str_abc_id_idx
+              spans: [/'ap-southeast-2'/'68088706-02c6-47d1-b993-a421cd761f2b' - /'ap-southeast-2'/'68088706-02c6-47d1-b993-a421cd761f2b']
+

pkg/sql/opt/distribution/distribution.go

@@ -155,25 +155,33 @@ func GetDEnumAsStringFromConstantExpr(expr opt.Expr) (enumAsString string, ok bo
 // of the input REGIONAL BY ROW table holding the crdb_region column, and
 // inputDistribution is the distribution of the operation on that table
 // (Scan or LocalityOptimizedSearch).
+// The distribution of the lookup join is returned, plus the first lookup index
+// column, as matched in the lookup with crdbRegionColID (if it can be
+// determined, otherwise zero).
 func BuildLookupJoinLookupTableDistribution(
 	ctx context.Context,
 	evalCtx *eval.Context,
 	lookupJoin *memo.LookupJoinExpr,
 	crdbRegionColID opt.ColumnID,
 	inputDistribution physical.Distribution,
-) (provided physical.Distribution) {
+) (firstLookupIndexCol opt.ColumnID, provided physical.Distribution) {
 	lookupTableMeta := lookupJoin.Memo().Metadata().TableMeta(lookupJoin.Table)
 	lookupTable := lookupTableMeta.Table
 
+	idx := lookupTable.Index(lookupJoin.Index)
+	col := idx.Column(0)
+	ord := col.Ordinal()
+	colIDOfFirstLookupIndexColumn := lookupTableMeta.MetaID.ColumnID(ord)
+
 	if lookupJoin.LocalityOptimized || lookupJoin.ChildOfLocalityOptimizedSearch {
 		provided.FromLocality(evalCtx.Locality)
-		return provided
+		return 0 /* firstLookupIndexCol */, provided
 	} else if lookupTable.IsGlobalTable() {
 		provided.FromLocality(evalCtx.Locality)
-		return provided
+		return 0 /* firstLookupIndexCol */, provided
 	} else if homeRegion, ok := lookupTable.HomeRegion(); ok {
 		provided.Regions = []string{homeRegion}
-		return provided
+		return 0 /* firstLookupIndexCol */, provided
 	} else if lookupTable.IsRegionalByRow() {
 		if len(lookupJoin.KeyCols) > 0 {
 			inputExpr := lookupJoin.Input
@@ -182,46 +190,46 @@ func BuildLookupJoinLookupTableDistribution(
 				if filterExpr, ok := invertedJoinExpr.GetConstExprFromFilter(firstKeyColID); ok {
 					if homeRegion, ok = GetDEnumAsStringFromConstantExpr(filterExpr); ok {
 						provided.Regions = []string{homeRegion}
-						return provided
+						return colIDOfFirstLookupIndexColumn, provided
 					}
 				}
 			} else if projectExpr, ok := inputExpr.(*memo.ProjectExpr); ok {
 				regionName := projectExpr.GetProjectedEnumConstant(firstKeyColID)
 				if regionName != "" {
 					provided.Regions = []string{regionName}
-					return provided
+					return colIDOfFirstLookupIndexColumn, provided
 				}
 			}
 			if crdbRegionColID == firstKeyColID {
 				provided.FromIndexScan(ctx, evalCtx, lookupTableMeta, lookupJoin.Index, nil)
 				if !inputDistribution.Any() &&
 					(provided.Any() || len(provided.Regions) > len(inputDistribution.Regions)) {
-					return inputDistribution
+					return colIDOfFirstLookupIndexColumn, inputDistribution
 				}
-				return provided
+				return 0 /* firstLookupIndexCol */, provided
 			}
 		} else if len(lookupJoin.LookupJoinPrivate.LookupExpr) > 0 {
 			if filterIdx, ok := lookupJoin.GetConstPrefixFilter(lookupJoin.Memo().Metadata()); ok {
 				firstIndexColEqExpr := lookupJoin.LookupJoinPrivate.LookupExpr[filterIdx].Condition
 				if firstIndexColEqExpr.Op() == opt.EqOp {
 					if regionName, ok := GetDEnumAsStringFromConstantExpr(firstIndexColEqExpr.Child(1)); ok {
 						provided.Regions = []string{regionName}
-						return provided
+						return colIDOfFirstLookupIndexColumn, provided
 					}
 				}
 			} else if lookupJoin.ColIsEquivalentWithLookupIndexPrefix(lookupJoin.Memo().Metadata(), crdbRegionColID) {
 				// We have a `crdb_region = crdb_region` term in `LookupJoinPrivate.LookupExpr`.
 				provided.FromIndexScan(ctx, evalCtx, lookupTableMeta, lookupJoin.Index, nil)
 				if !inputDistribution.Any() &&
 					(provided.Any() || len(provided.Regions) > len(inputDistribution.Regions)) {
-					return inputDistribution
+					return colIDOfFirstLookupIndexColumn, inputDistribution
 				}
-				return provided
+				return 0 /* firstLookupIndexCol */, provided
 			}
 		}
 	}
 	provided.FromIndexScan(ctx, evalCtx, lookupTableMeta, lookupJoin.Index, nil)
-	return provided
+	return 0 /* firstLookupIndexCol */, provided
 }
 
 // BuildInvertedJoinLookupTableDistribution builds the Distribution that results

pkg/sql/opt/xform/coster.go

@@ -1062,6 +1062,13 @@ func (c *coster) getCRBDRegionColFromInput(
 				return 0, physical.Distribution{}
 			}
 		}
+		if lookupJoinExpr, ok := maybeScan.(*memo.LookupJoinExpr); ok {
+			crdbRegionColID, inputDistribution = c.getCRBDRegionColFromInput(lookupJoinExpr, required)
+			crdbRegionColID, inputDistribution =
+				distribution.BuildLookupJoinLookupTableDistribution(
+					c.ctx, c.evalCtx, lookupJoinExpr, crdbRegionColID, inputDistribution)
+			return crdbRegionColID, inputDistribution
+		}
 		if localityOptimizedScan, ok := maybeScan.(*memo.LocalityOptimizedSearchExpr); ok {
 			maybeScan = localityOptimizedScan.Local
 			needRemap = true
@@ -1119,7 +1126,8 @@ func (c *coster) computeLookupJoinCost(
 		join.LocalityOptimized,
 	)
 	crdbRegionColID, inputDistribution := c.getCRBDRegionColFromInput(join, required)
-	provided := distribution.BuildLookupJoinLookupTableDistribution(c.ctx, c.evalCtx, join, crdbRegionColID, inputDistribution)
+	_, provided := distribution.BuildLookupJoinLookupTableDistribution(
+		c.ctx, c.evalCtx, join, crdbRegionColID, inputDistribution)
 	extraCost := c.distributionCost(provided)
 	cost += extraCost
 	return cost

pkg/sql/opt/xform/join_funcs.go

@@ -1890,7 +1890,7 @@ func (c *CustomFuncs) CanMaybeGenerateLocalityOptimizedSearchOfLookupJoins(
 // are done in the local region.
 func (c *CustomFuncs) LookupsAreLocal(lookupJoinExpr *memo.LookupJoinExpr) bool {
 	var inputDistribution physical.Distribution
-	provided := distribution.BuildLookupJoinLookupTableDistribution(c.e.ctx, c.e.f.EvalContext(), lookupJoinExpr, 0, inputDistribution)
+	_, provided := distribution.BuildLookupJoinLookupTableDistribution(c.e.ctx, c.e.f.EvalContext(), lookupJoinExpr, 0, inputDistribution)
 	if provided.Any() || len(provided.Regions) != 1 {
 		return false
 	}