Skip to content

Commit

Permalink
Merge #86631
Browse files Browse the repository at this point in the history 
86631: opt: don't propagate rule props when corresponding rule is disabled r=DrewKimball a=DrewKimball

This commit decreases the likelihood of rule cycles occuring during optimizer
tests with disabled rules. `DerivePruneCols` and `DeriveRejectNullCols`
now check whether propagating their corresponding properties would trigger
a disabled rule (if it wasn't disabled), and avoid propagating in that case.
This is necessary to avoid cases where a `Select` or `Project` gets repeatedly
pushed down and eliminated.

Addresses #86308

Release note: None

Release justification: testing-only change

Co-authored-by: DrewKimball <[email protected]>
  • Loading branch information
@DrewKimball
craig[bot] and DrewKimball committed Aug 30, 2022
2 parents d676dcb + 01c425e commit f200792
Show file tree
Hide file tree
Showing 7 changed files with 218 additions and 36 deletions.
13 changes: 13 additions & 0 deletions pkg/sql/opt/norm/factory.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/cockroachdb/cockroach/pkg/util"
"github.com/cockroachdb/cockroach/pkg/util/buildutil"
"github.com/cockroachdb/cockroach/pkg/util/errorutil"
"github.com/cockroachdb/errors"
Expand Down Expand Up @@ -95,6 +96,10 @@ type Factory struct {
// methods. It is incremented when a constructor function is called, and
// decremented when a constructor function returns.
constructorStackDepth int

// disabledRules is a set of rules that are not allowed to run, used when
// rules are disabled during testing to prevent rule cycles.
disabledRules util.FastIntSet
}

// maxConstructorStackDepth is the maximum allowed depth of a constructor call
Expand Down Expand Up @@ -195,6 +200,14 @@ func (f *Factory) NotifyOnAppliedRule(appliedRule AppliedRuleFunc) {
f.appliedRule = appliedRule
}

// SetDisabledRules is used to prevent normalization rule cycles when rules are
// disabled during testing. SetDisabledRules does not prevent rules from
// matching - rather, it notifies the Factory that rules have been prevented
// from matching using NotifyOnMatchedRule.
func (f *Factory) SetDisabledRules(disabledRules util.FastIntSet) {
f.disabledRules = disabledRules
}

// Memo returns the memo structure that the factory is operating upon.
func (f *Factory) Memo() *memo.Memo {
return f.mem
Expand Down
83 changes: 70 additions & 13 deletions pkg/sql/opt/norm/prune_cols_funcs.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/cockroachdb/cockroach/pkg/util"
)

// NeededGroupingCols returns the columns needed by a grouping operator's
Expand Down Expand Up @@ -251,7 +252,7 @@ func (c *CustomFuncs) NeededMutationFetchCols(
// needed columns from that. See the props.Relational.Rule.PruneCols comment for
// more details.
func (c *CustomFuncs) CanPruneCols(target memo.RelExpr, neededCols opt.ColSet) bool {
return !DerivePruneCols(target).SubsetOf(neededCols)
return !DerivePruneCols(target, c.f.disabledRules).SubsetOf(neededCols)
}

// CanPruneAggCols returns true if one or more of the target aggregations is not
Expand Down Expand Up @@ -309,7 +310,7 @@ func (c *CustomFuncs) PruneCols(target memo.RelExpr, neededCols opt.ColSet) memo
// Get the subset of the target expression's output columns that should
// not be pruned. Don't prune if the target output column is needed by a
// higher-level expression, or if it's not part of the PruneCols set.
pruneCols := DerivePruneCols(target).Difference(neededCols)
pruneCols := DerivePruneCols(target, c.f.disabledRules).Difference(neededCols)
colSet := c.OutputCols(target).Difference(pruneCols)
return c.f.ConstructProject(target, memo.EmptyProjectionsExpr, colSet)
}
Expand Down Expand Up @@ -499,7 +500,12 @@ func (c *CustomFuncs) PruneWindows(needed opt.ColSet, windows memo.WindowsExpr)
// what columns it allows to be pruned. Note that if an operator allows columns
// to be pruned, then there must be logic in the PruneCols method to actually
// prune those columns when requested.
func DerivePruneCols(e memo.RelExpr) opt.ColSet {
//
// disabledRules is the set of rules currently disabled, only used when rules
// are randomly disabled for testing. It is used to prevent propagating the
// PruneCols property when the corresponding column-pruning normalization rule
// is disabled. This prevents rule cycles during testing.
func DerivePruneCols(e memo.RelExpr, disabledRules util.FastIntSet) opt.ColSet {
relProps := e.Relational()
if relProps.IsAvailable(props.PruneCols) {
return relProps.Rule.PruneCols
Expand All @@ -508,33 +514,53 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {

switch e.Op() {
case opt.ScanOp, opt.ValuesOp, opt.WithScanOp:
if disabledRules.Contains(int(opt.PruneScanCols)) ||
disabledRules.Contains(int(opt.PruneValuesCols)) ||
disabledRules.Contains(int(opt.PruneWithScanCols)) {
// Avoid rule cycles.
break
}
// All columns can potentially be pruned from the Scan, Values, and WithScan
// operators.
relProps.Rule.PruneCols = relProps.OutputCols.Copy()

case opt.SelectOp:
if disabledRules.Contains(int(opt.PruneSelectCols)) {
// Avoid rule cycles.
break
}
// Any pruneable input columns can potentially be pruned, as long as they're
// not used by the filter.
sel := e.(*memo.SelectExpr)
relProps.Rule.PruneCols = DerivePruneCols(sel.Input).Copy()
relProps.Rule.PruneCols = DerivePruneCols(sel.Input, disabledRules).Copy()
usedCols := sel.Filters.OuterCols()
relProps.Rule.PruneCols.DifferenceWith(usedCols)

case opt.ProjectOp:
if disabledRules.Contains(int(opt.PruneProjectCols)) {
// Avoid rule cycles.
break
}
// All columns can potentially be pruned from the Project, if they're never
// used in a higher-level expression.
relProps.Rule.PruneCols = relProps.OutputCols.Copy()

case opt.InnerJoinOp, opt.LeftJoinOp, opt.RightJoinOp, opt.FullJoinOp,
opt.SemiJoinOp, opt.AntiJoinOp, opt.InnerJoinApplyOp, opt.LeftJoinApplyOp,
opt.SemiJoinApplyOp, opt.AntiJoinApplyOp:
if disabledRules.Contains(int(opt.PruneJoinLeftCols)) ||
disabledRules.Contains(int(opt.PruneJoinRightCols)) ||
disabledRules.Contains(int(opt.PruneSemiAntiJoinRightCols)) {
// Avoid rule cycles.
break
}
// Any pruneable columns from projected inputs can potentially be pruned, as
// long as they're not used by the right input (i.e. in Apply case) or by
// the join filter.
left := e.Child(0).(memo.RelExpr)
leftPruneCols := DerivePruneCols(left)
leftPruneCols := DerivePruneCols(left, disabledRules)
right := e.Child(1).(memo.RelExpr)
rightPruneCols := DerivePruneCols(right)
rightPruneCols := DerivePruneCols(right, disabledRules)

switch e.Op() {
case opt.SemiJoinOp, opt.SemiJoinApplyOp, opt.AntiJoinOp, opt.AntiJoinApplyOp:
Expand All @@ -548,6 +574,11 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {
relProps.Rule.PruneCols.DifferenceWith(onCols)

case opt.GroupByOp, opt.ScalarGroupByOp, opt.DistinctOnOp, opt.EnsureDistinctOnOp:
if disabledRules.Contains(int(opt.PruneGroupByCols)) ||
disabledRules.Contains(int(opt.PruneAggCols)) {
// Avoid rule cycles.
break
}
// Grouping columns can't be pruned, because they were used to group rows.
// However, aggregation columns can potentially be pruned.
groupingColSet := e.Private().(*memo.GroupingPrivate).GroupingCols
Expand All @@ -558,19 +589,28 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {
}

case opt.LimitOp, opt.OffsetOp:
if disabledRules.Contains(int(opt.PruneLimitCols)) ||
disabledRules.Contains(int(opt.PruneOffsetCols)) {
// Avoid rule cycles.
break
}
// Any pruneable input columns can potentially be pruned, as long as
// they're not used as an ordering column.
inputPruneCols := DerivePruneCols(e.Child(0).(memo.RelExpr))
inputPruneCols := DerivePruneCols(e.Child(0).(memo.RelExpr), disabledRules)
ordering := e.Private().(*props.OrderingChoice).ColSet()
relProps.Rule.PruneCols = inputPruneCols.Difference(ordering)

case opt.OrdinalityOp:
if disabledRules.Contains(int(opt.PruneOrdinalityCols)) {
// Avoid rule cycles.
break
}
// Any pruneable input columns can potentially be pruned, as long as
// they're not used as an ordering column. The new row number column
// cannot be pruned without adding an additional Project operator, so
// don't add it to the set.
ord := e.(*memo.OrdinalityExpr)
inputPruneCols := DerivePruneCols(ord.Input)
inputPruneCols := DerivePruneCols(ord.Input, disabledRules)
relProps.Rule.PruneCols = inputPruneCols.Difference(ord.Ordering.ColSet())

case opt.IndexJoinOp, opt.LookupJoinOp, opt.MergeJoinOp:
Expand All @@ -581,26 +621,39 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {
// currently a PruneCols rule for these operators.

case opt.ProjectSetOp:
if disabledRules.Contains(int(opt.PruneProjectSetCols)) {
// Avoid rule cycles.
break
}
// Any pruneable input columns can potentially be pruned, as long as
// they're not used in the Zip.
// TODO(rytaft): It may be possible to prune Zip columns, but we need to
// make sure that we still get the correct number of rows in the output.
projectSet := e.(*memo.ProjectSetExpr)
relProps.Rule.PruneCols = DerivePruneCols(projectSet.Input).Copy()
relProps.Rule.PruneCols = DerivePruneCols(projectSet.Input, disabledRules).Copy()
usedCols := projectSet.Zip.OuterCols()
relProps.Rule.PruneCols.DifferenceWith(usedCols)

case opt.UnionAllOp:
if disabledRules.Contains(int(opt.PruneUnionAllCols)) {
// Avoid rule cycles.
break
}
// Pruning can be beneficial as long as one of our inputs has advertised pruning,
// so that we can push down the project and eliminate the advertisement.
u := e.(*memo.UnionAllExpr)
pruneFromLeft := opt.TranslateColSet(DerivePruneCols(u.Left), u.LeftCols, u.OutCols)
pruneFromRight := opt.TranslateColSet(DerivePruneCols(u.Right), u.RightCols, u.OutCols)
pruneFromLeft := opt.TranslateColSet(DerivePruneCols(u.Left, disabledRules), u.LeftCols, u.OutCols)
pruneFromRight := opt.TranslateColSet(DerivePruneCols(u.Right, disabledRules), u.RightCols, u.OutCols)
relProps.Rule.PruneCols = pruneFromLeft.Union(pruneFromRight)

case opt.WindowOp:
if disabledRules.Contains(int(opt.PruneWindowInputCols)) ||
disabledRules.Contains(int(opt.PruneWindowOutputCols)) {
// Avoid rule cycles.
break
}
win := e.(*memo.WindowExpr)
relProps.Rule.PruneCols = DerivePruneCols(win.Input).Copy()
relProps.Rule.PruneCols = DerivePruneCols(win.Input, disabledRules).Copy()
relProps.Rule.PruneCols.DifferenceWith(win.Partition)
relProps.Rule.PruneCols.DifferenceWith(win.Ordering.ColSet())
for _, w := range win.Windows {
Expand All @@ -609,9 +662,13 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {
}

case opt.WithOp:
if disabledRules.Contains(int(opt.PruneWithCols)) {
// Avoid rule cycles.
break
}
// WithOp passes through its input unchanged, so it has the same pruning
// characteristics as its input.
relProps.Rule.PruneCols = DerivePruneCols(e.(*memo.WithExpr).Main)
relProps.Rule.PruneCols = DerivePruneCols(e.(*memo.WithExpr).Main, disabledRules)

default:
// Don't allow any columns to be pruned, since that would trigger the
Expand Down
59 changes: 47 additions & 12 deletions pkg/sql/opt/norm/reject_nulls_funcs.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,15 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/opt"
"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
"github.com/cockroachdb/cockroach/pkg/util"
"github.com/cockroachdb/errors"
)

// RejectNullCols returns the set of columns that are candidates for NULL
// rejection filter pushdown. See the Relational.Rule.RejectNullCols comment for
// more details.
func (c *CustomFuncs) RejectNullCols(in memo.RelExpr) opt.ColSet {
return DeriveRejectNullCols(in)
return DeriveRejectNullCols(in, c.f.disabledRules)
}

// HasNullRejectingFilter returns true if the filter causes some of the columns
Expand Down Expand Up @@ -118,7 +119,12 @@ func (c *CustomFuncs) NullRejectProjections(
// DeriveRejectNullCols returns the set of columns that are candidates for NULL
// rejection filter pushdown. See the Relational.Rule.RejectNullCols comment for
// more details.
func DeriveRejectNullCols(in memo.RelExpr) opt.ColSet {
//
// disabledRules is the set of rules currently disabled, only used when rules
// are randomly disabled for testing. It is used to prevent propagating the
// RejectNullCols property when the corresponding column-pruning normalization
// rule is disabled. This prevents rule cycles during testing.
func DeriveRejectNullCols(in memo.RelExpr, disabledRules util.FastIntSet) opt.ColSet {
// Lazily calculate and store the RejectNullCols value.
relProps := in.Relational()
if relProps.IsAvailable(props.RejectNullCols) {
Expand All @@ -131,37 +137,66 @@ func DeriveRejectNullCols(in memo.RelExpr) opt.ColSet {
case opt.InnerJoinOp, opt.InnerJoinApplyOp:
// Pass through null-rejecting columns from both inputs.
if in.Child(0).(memo.RelExpr).Relational().OuterCols.Empty() {
relProps.Rule.RejectNullCols.UnionWith(DeriveRejectNullCols(in.Child(0).(memo.RelExpr)))
relProps.Rule.RejectNullCols.UnionWith(
DeriveRejectNullCols(in.Child(0).(memo.RelExpr), disabledRules),
)
}
if in.Child(1).(memo.RelExpr).Relational().OuterCols.Empty() {
relProps.Rule.RejectNullCols.UnionWith(DeriveRejectNullCols(in.Child(1).(memo.RelExpr)))
relProps.Rule.RejectNullCols.UnionWith(
DeriveRejectNullCols(in.Child(1).(memo.RelExpr), disabledRules),
)
}

case opt.LeftJoinOp, opt.LeftJoinApplyOp:
if disabledRules.Contains(int(opt.RejectNullsLeftJoin)) {
// Avoid rule cycles.
break
}
// Pass through null-rejection columns from left input, and request null-
// rejection on right columns.
if in.Child(0).(memo.RelExpr).Relational().OuterCols.Empty() {
relProps.Rule.RejectNullCols.UnionWith(DeriveRejectNullCols(in.Child(0).(memo.RelExpr)))
relProps.Rule.RejectNullCols.UnionWith(
DeriveRejectNullCols(in.Child(0).(memo.RelExpr), disabledRules),
)
}
relProps.Rule.RejectNullCols.UnionWith(in.Child(1).(memo.RelExpr).Relational().OutputCols)

case opt.RightJoinOp:
if disabledRules.Contains(int(opt.RejectNullsRightJoin)) {
// Avoid rule cycles.
break
}
// Pass through null-rejection columns from right input, and request null-
// rejection on left columns.
relProps.Rule.RejectNullCols.UnionWith(in.Child(0).(memo.RelExpr).Relational().OutputCols)
if in.Child(1).(memo.RelExpr).Relational().OuterCols.Empty() {
relProps.Rule.RejectNullCols.UnionWith(DeriveRejectNullCols(in.Child(1).(memo.RelExpr)))
relProps.Rule.RejectNullCols.UnionWith(
DeriveRejectNullCols(in.Child(1).(memo.RelExpr), disabledRules),
)
}

case opt.FullJoinOp:
if disabledRules.Contains(int(opt.RejectNullsLeftJoin)) ||
disabledRules.Contains(int(opt.RejectNullsRightJoin)) {
// Avoid rule cycles.
break
}
// Request null-rejection on all output columns.
relProps.Rule.RejectNullCols.UnionWith(relProps.OutputCols)

case opt.GroupByOp, opt.ScalarGroupByOp:
relProps.Rule.RejectNullCols.UnionWith(deriveGroupByRejectNullCols(in))
if disabledRules.Contains(int(opt.RejectNullsGroupBy)) {
// Avoid rule cycles.
break
}
relProps.Rule.RejectNullCols.UnionWith(deriveGroupByRejectNullCols(in, disabledRules))

case opt.ProjectOp:
relProps.Rule.RejectNullCols.UnionWith(deriveProjectRejectNullCols(in))
if disabledRules.Contains(int(opt.RejectNullsProject)) {
// Avoid rule cycles.
break
}
relProps.Rule.RejectNullCols.UnionWith(deriveProjectRejectNullCols(in, disabledRules))

case opt.ScanOp:
relProps.Rule.RejectNullCols.UnionWith(deriveScanRejectNullCols(in))
Expand Down Expand Up @@ -191,7 +226,7 @@ func DeriveRejectNullCols(in memo.RelExpr) opt.ColSet {
// ignored because all rows in each group must have the same value for this
// column, so it doesn't matter which rows are filtered.
//
func deriveGroupByRejectNullCols(in memo.RelExpr) opt.ColSet {
func deriveGroupByRejectNullCols(in memo.RelExpr, disabledRules util.FastIntSet) opt.ColSet {
input := in.Child(0).(memo.RelExpr)
aggs := *in.Child(1).(*memo.AggregationsExpr)

Expand Down Expand Up @@ -222,7 +257,7 @@ func deriveGroupByRejectNullCols(in memo.RelExpr) opt.ColSet {
}
savedInColID = inColID

if !DeriveRejectNullCols(input).Contains(inColID) {
if !DeriveRejectNullCols(input, disabledRules).Contains(inColID) {
// Input has not requested null rejection on the input column.
return opt.ColSet{}
}
Expand Down Expand Up @@ -278,8 +313,8 @@ func (c *CustomFuncs) MakeNullRejectFilters(nullRejectCols opt.ColSet) memo.Filt
// child operator (for example, an outer join that may be simplified). This
// prevents filters from getting in the way of other rules.
//
func deriveProjectRejectNullCols(in memo.RelExpr) opt.ColSet {
rejectNullCols := DeriveRejectNullCols(in.Child(0).(memo.RelExpr))
func deriveProjectRejectNullCols(in memo.RelExpr, disabledRules util.FastIntSet) opt.ColSet {
rejectNullCols := DeriveRejectNullCols(in.Child(0).(memo.RelExpr), disabledRules)
projections := *in.Child(1).(*memo.ProjectionsExpr)
var projectionsRejectCols opt.ColSet

Expand Down
Loading

0 comments on commit f200792

Please sign in to comment.