sql/opt: handle non-unique parent for RBR region inference

pkg/sql/logictest/testdata/logic_test/regional_by_row_foreign_key

@@ -2440,3 +2440,54 @@ statement ok
 DROP TABLE parent_160013
 
 subtest end
+
+subtest non_unique_parent
+
+# Regression test for #170768: a REGIONAL BY ROW child can set
+# infer_rbr_region_col_using_constraint even when the parent's non-region
+# foreign-key columns aren't covered by a unique constraint. The optimizer
+# wraps the lookup join in a DistinctOn so each child input row maps to at
+# most one parent row; the chosen parent row's region is arbitrary if more
+# than one parent row matches.
+statement ok
+CREATE TABLE p_nonunique (
+  region crdb_internal_region NOT NULL,
+  k INT NOT NULL,
+  PRIMARY KEY (region, k)
+)
+
+statement ok
+INSERT INTO p_nonunique VALUES ('ca-central-1', 1), ('ap-southeast-2', 2), ('us-east-1', 3)
+
+statement ok
+CREATE TABLE c_nonunique (
+  pk INT PRIMARY KEY,
+  data INT,
+  CONSTRAINT fk_region FOREIGN KEY (crdb_region, pk) REFERENCES p_nonunique (region, k)
+) WITH (infer_rbr_region_col_using_constraint = 'fk_region') LOCALITY REGIONAL BY ROW
+
+# Before #170768's fix the INSERT panicked with an assertion failure
+# ("lookup columns using constraint X must be a lax key"). It now succeeds:
+# the single matching parent row's region is selected.
+statement ok
+INSERT INTO c_nonunique (pk, data) VALUES (1, 100)
+
+query IT
+SELECT pk, crdb_region FROM c_nonunique
+----
+1  ca-central-1
+
+# With two parent rows sharing the join key, the inference still completes
+# (picking one arbitrarily) rather than failing.
+statement ok
+INSERT INTO p_nonunique VALUES ('ca-central-1', 10), ('us-east-1', 10)
+
+statement ok
+INSERT INTO c_nonunique (pk, data) VALUES (10, 100)
+
+query I
+SELECT count(*) FROM c_nonunique WHERE pk = 10
+----
+1
+
+subtest end

pkg/sql/opt/norm/decorrelate_funcs.go

@@ -518,81 +518,6 @@ func (c *CustomFuncs) ConstructApplyJoin(
 	panic(errors.AssertionFailedf("unexpected join operator: %v", redact.Safe(joinOp)))
 }
 
-// EnsureKey finds the shortest strong key for the input expression. If no
-// strong key exists and the input expression is a Scan or a Scan wrapped in a
-// Select, EnsureKey returns a new Scan (possibly wrapped in a Select) with the
-// preexisting primary key for the table. If the input is not a Scan or
-// Select(Scan), EnsureKey wraps the input in an Ordinality operator, which
-// provides a key column by uniquely numbering the rows. EnsureKey returns the
-// input expression (perhaps augmented with a key column(s) or wrapped by
-// Ordinality).
-func (c *CustomFuncs) EnsureKey(in memo.RelExpr) memo.RelExpr {
-	// Try to add the preexisting primary key if the input is a Scan or Scan
-	// wrapped in a Select.
-	if res, ok := c.tryFindExistingKey(in); ok {
-		return res
-	}
-
-	// Otherwise, wrap the input in an Ordinality operator.
-	colID := c.f.Metadata().AddColumn("rownum", types.Int)
-	private := memo.OrdinalityPrivate{ColID: colID, ForDuplicateRemoval: true}
-	return c.f.ConstructOrdinality(in, &private)
-}
-
-// 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
-		}
-
-	case *memo.SelectExpr:
-		input, foundKey := c.tryFindExistingKey(t.Input)
-		if foundKey {
-			return c.f.ConstructSelect(input, t.Filters), true
-		}
-	}
-
-	return nil, false
-}
-
-// KeyCols returns a column set consisting of the columns that make up the
-// candidate key for the input expression (a key must be present).
-func (c *CustomFuncs) KeyCols(in memo.RelExpr) opt.ColSet {
-	keyCols, ok := c.CandidateKey(in)
-	if !ok {
-		panic(errors.AssertionFailedf("expected expression to have key"))
-	}
-	return keyCols
-}
-
-// NonKeyCols returns a column set consisting of the output columns of the given
-// input, minus the columns that make up its candidate key (which it must have).
-func (c *CustomFuncs) NonKeyCols(in memo.RelExpr) opt.ColSet {
-	keyCols, ok := c.CandidateKey(in)
-	if !ok {
-		panic(errors.AssertionFailedf("expected expression to have key"))
-	}
-	return c.OutputCols(in).Difference(keyCols)
-}
-
 // MakeAggCols2 is similar to MakeAggCols, except that it allows two different
 // sets of aggregate functions to be added to the resulting Aggregations
 // operator, with one set appended to the other, like this:

pkg/sql/opt/norm/general_funcs.go

@@ -611,6 +611,81 @@ func (c *CustomFuncs) PrimaryKeyCols(table opt.TableID) opt.ColSet {
 	return tabMeta.IndexKeyColumns(cat.PrimaryIndex)
 }
 
+// KeyCols returns a column set consisting of the columns that make up the
+// candidate key for the input expression (a key must be present).
+func (c *CustomFuncs) KeyCols(in memo.RelExpr) opt.ColSet {
+	keyCols, ok := c.CandidateKey(in)
+	if !ok {
+		panic(errors.AssertionFailedf("expected expression to have key"))
+	}
+	return keyCols
+}
+
+// NonKeyCols returns a column set consisting of the output columns of the given
+// input, minus the columns that make up its candidate key (which it must have).
+func (c *CustomFuncs) NonKeyCols(in memo.RelExpr) opt.ColSet {
+	keyCols, ok := c.CandidateKey(in)
+	if !ok {
+		panic(errors.AssertionFailedf("expected expression to have key"))
+	}
+	return c.OutputCols(in).Difference(keyCols)
+}
+
+// EnsureKey finds the shortest strong key for the input expression. If no
+// strong key exists and the input expression is a Scan or a Scan wrapped in a
+// Select, EnsureKey returns a new Scan (possibly wrapped in a Select) with the
+// preexisting primary key for the table. If the input is not a Scan or
+// Select(Scan), EnsureKey wraps the input in an Ordinality operator, which
+// provides a key column by uniquely numbering the rows. EnsureKey returns the
+// input expression (perhaps augmented with a key column(s) or wrapped by
+// Ordinality).
+func (c *CustomFuncs) EnsureKey(in memo.RelExpr) memo.RelExpr {
+	// Try to add the preexisting primary key if the input is a Scan or Scan
+	// wrapped in a Select.
+	if res, ok := c.tryFindExistingKey(in); ok {
+		return res
+	}
+
+	// Otherwise, wrap the input in an Ordinality operator.
+	colID := c.f.Metadata().AddColumn("rownum", types.Int)
+	private := memo.OrdinalityPrivate{ColID: colID, ForDuplicateRemoval: true}
+	return c.f.ConstructOrdinality(in, &private)
+}
+
+// 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
+		}
+
+	case *memo.SelectExpr:
+		input, foundKey := c.tryFindExistingKey(t.Input)
+		if foundKey {
+			return c.f.ConstructSelect(input, t.Filters), true
+		}
+	}
+
+	return nil, false
+}
+
 // ----------------------------------------------------------------------
 //
 // Property functions

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -1048,20 +1048,40 @@ func (mb *mutationBuilder) maybeAddRegionColLookup(op opt.Operator) {
 		// The scan is exempt from RLS to maintain data integrity.
 		cat.PolicyScopeExempt,
 	)
-	if !refScope.expr.Relational().FuncDeps.ColsAreLaxKey(refLookupCols) {
-		// The lookup columns must be a lax key, otherwise the join may return
-		// multiple rows for a single row in the target table. This should already
-		// be enforced by the foreign-key constraint.
-		panic(errors.AssertionFailedf(
-			"lookup columns using constraint %q must be a lax key", lookupFK.Name()))
-	}
 	var joinFlags memo.JoinFlags
 	if mb.b.evalCtx.SessionData().PreferLookupJoinsForFKs {
 		joinFlags = memo.PreferLookupJoinIntoRight
 	}
+	parentNonKey := !refScope.expr.Relational().FuncDeps.ColsAreLaxKey(refLookupCols)
+	if parentNonKey {
+		// The non-region foreign-key columns aren't a lax key on the parent, so
+		// the join could return multiple rows per child input row. We will use an
+		// unordered DistinctOn to select an arbitrary matching parent row for each
+		// child row. This requires a key on the child rows, so ensure it here.
+		mb.outScope.expr = mb.b.factory.CustomFuncs().EnsureKey(mb.outScope.expr)
+	}
+	inputCols := mb.outScope.expr.Relational().OutputCols
 	mb.outScope.expr = mb.b.factory.ConstructLeftJoin(
 		mb.outScope.expr, refScope.expr, joinCond, &memo.JoinPrivate{Flags: joinFlags},
 	)
+	if parentNonKey {
+		// Use an unordered DistinctOn to arbitrarily select matching parent rows.
+		// Group on all input columns from before the join (they form a superkey
+		// after EnsureKey); optimizer rules can simplify the grouping columns
+		// using the input key. The parent's region column is the only parent
+		// column needed downstream, and FirstAgg picks one arbitrary value from
+		// the matching parent rows.
+		aggs := memo.AggregationsExpr{
+			f.ConstructAggregationsItem(
+				f.ConstructFirstAgg(f.ConstructVariable(lookupRegionColID)),
+				lookupRegionColID,
+			),
+		}
+		groupingPrivate := memo.GroupingPrivate{GroupingCols: inputCols}
+		mb.outScope.expr = f.ConstructDistinctOn(
+			mb.outScope.expr, aggs, &groupingPrivate,
+		)
+	}
 	// Build a CASE expression to determine the final value of the region column.
 	// Use the looked-up value if non-NULL, and otherwise use the default value
 	// which was already projected in the input.