Merge cockroachdb#130978

130978: opt: fix SplitDisjunctionOfJoinTerms and SplitDisjunctionOfAntiJoinTerms r=mgartner a=mgartner #### opt: fix test assertion message A test-only assertion with the message "projection passes through columns not in input" has been fixed to correctly print the columns that are passthrough columns and not in the input. Previously, it incorrectly printed input columns that were not passthrough columns. Release note: None #### opt: fix SplitDisjunctionOfJoinTerms and SplitDisjunctionOfAntiJoinTerms The exploration rules SplitDisjunctionOfJoinTerms and SplitDisjunctionOfAntiJoinTerms previously constructed expressions during the match phase of the rule. This is an anti-pattern. The match logic occurs before the `matchedRule` callback is invoked, which is used to disable exploration rules under the `norm` test directive. As a result these rules could perform some constructions of new expressions in `norm` tests where these exploration rules should be disabled. This confused me quite a bit while debugging a problem with the rules. The rules have been updated so that the eligibility of the rules are checked during the match phase, and new expressions are only constructed during the replace phase. Epic: None Release note: None #### opt: optimize SplitDisjunctionOfJoinTerms and SplitDisjunctionOfAntiJoinTerms Duplicate computation that occurred in both the match and replace phases of the SplitDisjunctionOfJoinTerms and SplitDisjunctionOfAntiJoinTerms rules has been eliminated by performing the computation in only the match phase, and passing the results to a custom function in the replace phase. Release note: None Co-authored-by: Marcus Gartner <[email protected]>
cthumuluru-crdb · Sep 18, 2024 · a818801 · a818801
2 parents 93f325f + 02c614c
commit a818801
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 41 deletions.
diff --git a/pkg/sql/opt/memo/check_expr.go b/pkg/sql/opt/memo/check_expr.go
@@ -93,7 +93,7 @@ func (m *Memo) CheckExpr(e opt.Expr) {
 		if !t.Passthrough.SubsetOf(t.Input.Relational().OutputCols) {
 			panic(errors.AssertionFailedf(
 				"projection passes through columns not in input: %v",
-				t.Input.Relational().OutputCols.Difference(t.Passthrough),
+				t.Passthrough.Difference(t.Input.Relational().OutputCols),
 			))
 		}
 		for _, item := range t.Projections {

diff --git a/pkg/sql/opt/xform/join_funcs.go b/pkg/sql/opt/xform/join_funcs.go
@@ -1536,6 +1536,35 @@ func (c *CustomFuncs) makeFilteredSelectForJoin(
 	return newSelect
 }
 
+// CanSplitJoinWithDisjuncts returns true if the given join can be split into a
+// union of two joins. See SplitJoinWithDisjuncts for more details.
+func (c *CustomFuncs) CanSplitJoinWithDisjuncts(
+	joinRel memo.RelExpr, joinFilters memo.FiltersExpr,
+) (firstOnClause, secondOnClause opt.ScalarExpr, itemToReplace *memo.FiltersItem, ok bool) {
+	leftInput := joinRel.Child(0).(memo.RelExpr)
+	rightInput := joinRel.Child(1).(memo.RelExpr)
+
+	switch joinRel.Op() {
+	case opt.InnerJoinOp, opt.SemiJoinOp, opt.AntiJoinOp:
+		// Do nothing
+	default:
+		panic(errors.AssertionFailedf("expected joinRel to be inner, semi, or anti-join"))
+	}
+
+	origLeftScan, _, ok := c.getfilteredCanonicalScan(leftInput)
+	if !ok {
+		return nil, nil, nil, false
+	}
+
+	origRightScan, _, ok := c.getfilteredCanonicalScan(rightInput)
+	if !ok {
+		return nil, nil, nil, false
+	}
+
+	// Look for a disjunction of equijoin predicates.
+	return c.splitDisjunctionForJoin(joinRel, joinFilters, origLeftScan, origRightScan)
+}
+
 // SplitJoinWithDisjuncts checks a join relation for a disjunction of predicates
 // in an InnerJoin, SemiJoin or AntiJoin. If present, and the inputs to the join
 // are canonical scans, or Selects from canonical scans, it builds two new join
@@ -1555,27 +1584,20 @@ func (c *CustomFuncs) makeFilteredSelectForJoin(
 // If there is no disjunction of predicates, or the join type is not one of the
 // supported join types listed above, ok=false is returned.
 func (c *CustomFuncs) SplitJoinWithDisjuncts(
-	joinRel memo.RelExpr, joinFilters memo.FiltersExpr,
+	joinRel memo.RelExpr,
+	joinFilters memo.FiltersExpr,
+	firstOnClause, secondOnClause opt.ScalarExpr,
+	itemToReplace *memo.FiltersItem,
 ) (
 	firstJoin memo.RelExpr,
 	secondJoin memo.RelExpr,
 	newRelationCols opt.ColSet,
 	aggCols opt.ColSet,
 	groupingCols opt.ColSet,
-	ok bool,
 ) {
-	notOkSplitJoin := func() (memo.RelExpr, memo.RelExpr, opt.ColSet, opt.ColSet, opt.ColSet, bool) {
-		emptyColSet := opt.ColSet{}
-		return nil, nil, emptyColSet, emptyColSet, emptyColSet, false
-	}
-
-	var joinPrivate *memo.JoinPrivate
-	var leftInput memo.RelExpr
-	var rightInput memo.RelExpr
-
-	joinPrivate = joinRel.Private().(*memo.JoinPrivate)
-	leftInput = joinRel.Child(0).(memo.RelExpr)
-	rightInput = joinRel.Child(1).(memo.RelExpr)
+	joinPrivate := joinRel.Private().(*memo.JoinPrivate)
+	leftInput := joinRel.Child(0).(memo.RelExpr)
+	rightInput := joinRel.Child(1).(memo.RelExpr)
 
 	switch joinRel.Op() {
 	case opt.InnerJoinOp, opt.SemiJoinOp, opt.AntiJoinOp:
@@ -1586,22 +1608,15 @@ func (c *CustomFuncs) SplitJoinWithDisjuncts(
 
 	origLeftScan, leftFilters, ok := c.getfilteredCanonicalScan(leftInput)
 	if !ok {
-		return notOkSplitJoin()
+		panic(errors.AssertionFailedf("expected join left input to have canonical scan"))
 	}
 	origLeftScanPrivate := &origLeftScan.ScanPrivate
 	origRightScan, rightFilters, ok := c.getfilteredCanonicalScan(rightInput)
 	if !ok {
-		return notOkSplitJoin()
+		panic(errors.AssertionFailedf("expected join right input to have canonical scan"))
 	}
 	origRightScanPrivate := &origRightScan.ScanPrivate
 
-	// Look for a disjunction of equijoin predicates.
-	firstOnClause, secondOnClause, itemToReplace, ok :=
-		c.splitDisjunctionForJoin(joinRel, joinFilters, origLeftScan, origRightScan)
-	if !ok {
-		return notOkSplitJoin()
-	}
-
 	// Add in the primary key columns so the caller can group by them to
 	// deduplicate results.
 	var leftSP *memo.ScanPrivate
@@ -1736,7 +1751,7 @@ func (c *CustomFuncs) SplitJoinWithDisjuncts(
 		panic(errors.AssertionFailedf("Unexpected join type while splitting disjuncted join predicates: %v",
 			joinRel.Op()))
 	}
-	return firstJoin, secondJoin, newRelationCols, aggCols, groupingCols, true
+	return firstJoin, secondJoin, newRelationCols, aggCols, groupingCols
 }
 
 // getfilteredCanonicalScan looks at a *ScanExpr or *SelectExpr "relation" and

diff --git a/pkg/sql/opt/xform/rules/join.opt b/pkg/sql/opt/xform/rules/join.opt
@@ -152,13 +152,11 @@
     $on:* &
         (Let
             (
-                $firstJoin
-                $secondJoin
-                $newRelationCols
-                $aggCols
-                $groupingCols
+                $firstOnClause
+                $secondOnClause
+                $itemToReplace
                 $ok
-            ):(SplitJoinWithDisjuncts (Root) $on)
+            ):(CanSplitJoinWithDisjuncts (Root) $on)
             $ok
         )
     *
@@ -167,7 +165,22 @@
 (Project
     (DistinctOn
         (UnionAll
-            $firstJoin
+            (Let
+                (
+                    $firstJoin
+                    $secondJoin
+                    $newRelationCols
+                    $aggCols
+                    $groupingCols
+                ):(SplitJoinWithDisjuncts
+                    (Root)
+                    $on
+                    $firstOnClause
+                    $secondOnClause
+                    $itemToReplace
+                )
+                $firstJoin
+            )
             $secondJoin
             (MakeSetPrivate
                 $firstJoinCols:(OutputCols $firstJoin)
@@ -195,22 +208,34 @@
     $on:* &
         (Let
             (
-                $firstJoin
-                $secondJoin
-                $newRelationCols
-                $aggCols
-                $groupingCols
+                $firstOnClause
+                $secondOnClause
+                $itemToReplace
                 $ok
-            ):(SplitJoinWithDisjuncts (Root) $on)
+            ):(CanSplitJoinWithDisjuncts (Root) $on)
             $ok
         )
-    *
 )
 =>
 (Project
     (DistinctOn
         (IntersectAll
-            $firstJoin
+            (Let
+                (
+                    $firstJoin
+                    $secondJoin
+                    $newRelationCols
+                    $aggCols
+                    $groupingCols
+                ):(SplitJoinWithDisjuncts
+                    (Root)
+                    $on
+                    $firstOnClause
+                    $secondOnClause
+                    $itemToReplace
+                )
+                $firstJoin
+            )
             $secondJoin
             (MakeSetPrivate
                 $firstJoinCols:(OutputCols $firstJoin)

diff --git a/pkg/sql/opt/xform/testdata/rules/disjunction_in_join b/pkg/sql/opt/xform/testdata/rules/disjunction_in_join
@@ -1,4 +1,4 @@
-# Test ORed ON clause predicates which may be split into unions or 
+# Test ORed ON clause predicates which may be split into unions or
 # intersections. The rewrite is cost-based, so may not always show up in the
 # query plan.
 # Use tables both with and without a primary key, to test when PK columns are