Skip to content

Commit

Permalink
opt: don't propagate rule props when corresponding rule is disabled
Browse files Browse the repository at this point in the history
This commit decreases the likelihood of rule cycles occuring during optimizer
tests with randomly 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
  • Loading branch information
DrewKimball committed Aug 30, 2022
1 parent 004533f commit 01c425e
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
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
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
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 01c425e

Please sign in to comment.