Skip to content

optimizer: Make EQProp call MSE::reduce only after a change in reduce_expr #31979

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggevay merged 2 commits into from
Mar 21, 2025
Merged

optimizer: Make EQProp call MSE::reduce only after a change in reduce_expr #31979

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e9e4f10
optimizer: Make EQProp call MSE::reduce only after a change in reduce…
ggevay Mar 21, 2025
3b3fdea
optimizer: add enable_less_reduce_in_eqprop feature flag
ggevay Mar 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/repr/src/optimize.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,8 @@ optimizer_feature_flags!({
enable_projection_pushdown_after_relation_cse: bool,
// See the feature flag of the same name.
enable_let_prefix_extraction: bool,
// See the feature flag of the same name.
enable_less_reduce_in_eqprop: bool,
});

/// A trait used to implement layered config construction.
Expand Down
2 changes: 2 additions & 0 deletions src/sql/src/plan/statement/ddl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4654,7 +4654,9 @@ pub fn unplan_create_cluster(
enable_join_prioritize_arranged,
enable_projection_pushdown_after_relation_cse,
enable_let_prefix_extraction: _,
enable_less_reduce_in_eqprop: _,
} = optimizer_feature_overrides;
// The ones from above that don't occur below are not wired up to cluster features.
let features_extracted = ClusterFeatureExtracted {
// Seen is ignored when unplanning.
seen: Default::default(),
Expand Down
1 change: 1 addition & 0 deletions src/sql/src/plan/statement/dml.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -435,6 +435,7 @@ impl TryFrom<ExplainPlanOptionExtracted> for ExplainConfig {
enable_projection_pushdown_after_relation_cse: v
.enable_projection_pushdown_after_relation_cse,
enable_let_prefix_extraction: Default::default(),
enable_less_reduce_in_eqprop: Default::default(),
},
})
}
Expand Down
7 changes: 7 additions & 0 deletions src/sql/src/session/vars/definitions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2209,6 +2209,12 @@ feature_flags!(
default: true,
enable_for_item_parsing: false,
},
{
name: enable_less_reduce_in_eqprop,
desc: "Run MSE::reduce in EquivalencePropagation only if reduce_expr changed something.",
default: true,
enable_for_item_parsing: false,
},
);

impl From<&super::SystemVars> for OptimizerFeatures {
Expand All @@ -2228,6 +2234,7 @@ impl From<&super::SystemVars> for OptimizerFeatures {
enable_projection_pushdown_after_relation_cse: vars
.enable_projection_pushdown_after_relation_cse(),
enable_let_prefix_extraction: vars.enable_let_prefix_extraction(),
enable_less_reduce_in_eqprop: vars.enable_less_reduce_in_eqprop(),
}
}
}
89 changes: 69 additions & 20 deletions src/transform/src/equivalence_propagation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ impl crate::Transform for EquivalencePropagation {
derived,
EquivalenceClasses::default(),
&mut get_equivalences,
ctx,
);

// Trace the plan as the result of `equivalence_propagation` before potentially applying
Expand Down Expand Up @@ -116,6 +117,7 @@ impl EquivalencePropagation {
derived: DerivedView,
mut outer_equivalences: EquivalenceClasses,
get_equivalences: &mut BTreeMap<Id, EquivalenceClasses>,
ctx: &mut TransformCtx,
) {
// TODO: The top-down traversal can be coded as a worklist, with arguments tupled and enqueued.
// This has the potential to do a lot more cloning (of `outer_equivalences`), and some care is needed
Expand Down Expand Up @@ -194,22 +196,40 @@ impl EquivalencePropagation {
let body = children_rev.next().unwrap();
let value = children_rev.next().unwrap();

self.apply(body.0, body.1, outer_equivalences.clone(), get_equivalences);
self.apply(
body.0,
body.1,
outer_equivalences.clone(),
get_equivalences,
ctx,
);

// We expect to find `id` in `get_equivalences`, as otherwise the binding is
// not referenced and can be removed.
if let Some(equivalences) = get_equivalences.get(&Id::Local(id)) {
self.apply(value.0, value.1, equivalences.clone(), get_equivalences);
self.apply(
value.0,
value.1,
equivalences.clone(),
get_equivalences,
ctx,
);
}
}
MirRelationExpr::LetRec { .. } => {
let mut child_iter = expr.children_mut().rev().zip(derived.children_rev());
// Continue in `body` with the outer equivalences.
let (body, view) = child_iter.next().unwrap();
self.apply(body, view, outer_equivalences, get_equivalences);
self.apply(body, view, outer_equivalences, get_equivalences, ctx);
// Continue recursively, but without the outer equivalences supplied to `body`.
for (child, view) in child_iter {
self.apply(child, view, EquivalenceClasses::default(), get_equivalences);
self.apply(
child,
view,
EquivalenceClasses::default(),
get_equivalences,
ctx,
);
}
}
MirRelationExpr::Project { input, outputs } => {
Expand All @@ -220,6 +240,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
MirRelationExpr::Map { input, scalars } => {
Expand All @@ -241,8 +262,10 @@ impl EquivalencePropagation {
let input_arity = input_types.len() - scalars.len();
for (index, expr) in scalars.iter_mut().enumerate() {
let reducer = input_equivalences.reducer();
reducer.reduce_expr(expr);
expr.reduce(&input_types[..(input_arity + index)]);
let changed = reducer.reduce_expr(expr);
if changed || !ctx.features.enable_less_reduce_in_eqprop {
expr.reduce(&input_types[..(input_arity + index)]);
}
// Introduce the fact relating the mapped expression and corresponding column.
// This allows subsequent expressions to be optimized with this information.
input_equivalences.classes.push(vec![
Expand All @@ -261,6 +284,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
}
Expand All @@ -278,8 +302,10 @@ impl EquivalencePropagation {
.expect("RelationType required");
let reducer = input_equivalences.reducer();
for expr in exprs.iter_mut() {
reducer.reduce_expr(expr);
expr.reduce(input_types.as_ref().unwrap());
let changed = reducer.reduce_expr(expr);
if changed || !ctx.features.enable_less_reduce_in_eqprop {
expr.reduce(input_types.as_ref().unwrap());
}
}
let input_arity = *derived
.last_child()
Expand All @@ -291,6 +317,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
}
Expand All @@ -309,8 +336,10 @@ impl EquivalencePropagation {
.expect("RelationType required");
let reducer = input_equivalences.reducer();
for expr in predicates.iter_mut() {
reducer.reduce_expr(expr);
expr.reduce(input_types.as_ref().unwrap());
let changed = reducer.reduce_expr(expr);
if changed || !ctx.features.enable_less_reduce_in_eqprop {
expr.reduce(input_types.as_ref().unwrap());
}
}
// Incorporate `predicates` into `outer_equivalences`.
let mut class = predicates.clone();
Expand All @@ -325,6 +354,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
}
Expand Down Expand Up @@ -403,9 +433,11 @@ impl EquivalencePropagation {
// Semijoin elimination currently fails if you do more advanced simplification than
// literal substitution.
let old = expr.clone();
reducer.reduce_expr(expr);
let changed = reducer.reduce_expr(expr);
let acceptable_sub = literal_domination(&old, expr);
expr.reduce(input_types.as_ref().unwrap());
if changed || !ctx.features.enable_less_reduce_in_eqprop {
expr.reduce(input_types.as_ref().unwrap());
}
if !acceptable_sub && !literal_domination(&old, expr) {
expr.clone_from(&old);
}
Expand Down Expand Up @@ -440,7 +472,7 @@ impl EquivalencePropagation {
}
}
push_equivalences.project(columns..(columns + child_arity));
self.apply(expr, child, push_equivalences, get_equivalences);
self.apply(expr, child, push_equivalences, get_equivalences, ctx);

columns += child_arity;
}
Expand Down Expand Up @@ -472,16 +504,20 @@ impl EquivalencePropagation {
// Semijoin elimination currently fails if you do more advanced simplification than
// literal substitution.
let old_key = key.clone();
reducer.reduce_expr(key);
let changed = reducer.reduce_expr(key);
let acceptable_sub = literal_domination(&old_key, key);
key.reduce(input_type.as_ref().unwrap());
if changed || !ctx.features.enable_less_reduce_in_eqprop {
key.reduce(input_type.as_ref().unwrap());
}
if !acceptable_sub && !literal_domination(&old_key, key) {
key.clone_from(&old_key);
}
}
for aggr in aggregates.iter_mut() {
reducer.reduce_expr(&mut aggr.expr);
aggr.expr.reduce(input_type.as_ref().unwrap());
let changed = reducer.reduce_expr(&mut aggr.expr);
if changed || !ctx.features.enable_less_reduce_in_eqprop {
aggr.expr.reduce(input_type.as_ref().unwrap());
}
// A count expression over a non-null expression can discard the expression.
if aggr.func == mz_expr::AggregateFunc::Count && !aggr.distinct {
let mut probe = aggr.expr.clone().call_is_null();
Expand Down Expand Up @@ -517,6 +553,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
MirRelationExpr::TopK {
Expand All @@ -540,9 +577,11 @@ impl EquivalencePropagation {
let reducer = input_equivalences.reducer();
if let Some(expr) = limit {
let old_expr = expr.clone();
reducer.reduce_expr(expr);
let changed = reducer.reduce_expr(expr);
let acceptable_sub = literal_domination(&old_expr, expr);
expr.reduce(input_types.as_ref().unwrap());
if changed || !ctx.features.enable_less_reduce_in_eqprop {
expr.reduce(input_types.as_ref().unwrap());
}
if !acceptable_sub && !literal_domination(&old_expr, expr) {
expr.clone_from(&old_expr);
}
Expand All @@ -557,6 +596,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
MirRelationExpr::Negate { input } => {
Expand All @@ -565,6 +605,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
MirRelationExpr::Threshold { input } => {
Expand All @@ -573,11 +614,18 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
MirRelationExpr::Union { .. } => {
for (child, derived) in expr.children_mut().rev().zip(derived.children_rev()) {
self.apply(child, derived, outer_equivalences.clone(), get_equivalences);
self.apply(
child,
derived,
outer_equivalences.clone(),
get_equivalences,
ctx,
);
}
}
MirRelationExpr::ArrangeBy { input, .. } => {
Expand All @@ -587,6 +635,7 @@ impl EquivalencePropagation {
derived.last_child(),
outer_equivalences,
get_equivalences,
ctx,
);
}
}
Expand Down