PoC: Explicit Invariant Regions in the Physical Optimizer

[2010YOUY01]

Which issue does this PR close?

• Closes #.

Rationale for this change

The default physical optimizer pipeline works for the optimizations implemented in DataFusion core. However, there are many hidden assumptions between optimizer rules that make the pipeline work as a whole (e.g., rule A must run before rule B to produce a valid plan).

These implicit constraints make extensions harder, whether adding new rules within DataFusion core or implementing custom pipelines downstream.

This PR introduces explicit invariant regions to make these hidden constraints visible and verifiable, making the optimizer easier to understand, extend, and maintain.

Note that this PR does not introduce new constraints. I think of this primarily as a "removing uncertainty" step. The assumptions already exist; this PR provides a mechanism to surface them so that:

1. Developers can understand the expectations between optimizer rules.
2. We can automatically test whether rules preserve those expectations.
3. Extensions can make more informed decisions about where to insert or rearrange rules.

The tradeoff is a certain learning curve for this mechanism. After prototyping the idea, I feel that once we become familiar with the concept, it can make long-term maintenance easier. However, it definitely looks confusing at first glance, and I'm still thinking about how to improve this part.

(Note: The implementation is currently underdocumented; most of the ideas are explained in this writeup.)

Two Types of Constraints

This section explains two types of constraints, with one example each. The following figure illustrates two invariant regions in an example optimizer pipeline, which will be explained further below.

                           Invariant Regions
Rule List                   Aggr              Req
--------------------------------------------------------
InitialPhysicalPlanner
                            +-- BEGIN
...                         |
LimitedDistinctAggregation  |
...                         |
                            +-- END
EnsureRequirements
                                              +-- BEGIN
...                                           |
                                              +-- END
ExtensionRule1                                (inactive)
ExtensionRule2                                (inactive)
EnsureRequirements                            (restore)
                                              +-- BEGIN
...                                           |
SanityCheckPlan                               |
                                              +-- END

• Region 1, CanonicalAggregatePlan (Aggr): Maintains canonical aggregation shapes in this region, the rules inside this region can kept simpler by only handling the canonical aggregate shape
• Region 2, InputRequirementsSatisfied (Req): Ensures execution requirements remain satisfied. The example demonstrated temporary turning off/on in the middle to make extension easier.

Type I: Optimization-Coverage Constraints (Soft)

Progressive refinement in the optimizer pipeline

The default optimizer pipeline is a progressive refinement process, where the possible plan shapes vary across stages.

For example:

• At rule 1, a hash join might have only one canonical representation.
• At rule 20, the same logical join might have five possible physical representations.

A rule that rewrites joins could be much simpler when placed at an earlier stage.

Requiring every rule to handle every possible representation, regardless of its position, can unnecessarily bloat the implementation of earlier rules.

Making soft constraints explicit

Built-in rules already rely on such canonicalization assumptions, allowing them to match only canonical plan shapes and keep their implementations simple.

Soft constraints provide a mechanism to make these assumptions explicit and verifiable, so custom optimizer pipelines can either:

• Follow the constraints to preserve the same optimization coverage. (Another interpretation is that these constraints help extension rules find suitable positions in the pipeline.)
• Opt out of the constraints, accepting the possibility of losing optimization coverage from built-in rules.

Example: LimitedDistinctAggregation

See #26069 for a concrete example of how relying on a canonical aggregation shape can simplify an optimizer rule.

The corresponding invariant region would be:

CanonicalAggregatePlan
- Established: After InitialPhysicalPlanner
- Ends: After LimitedDistinctAggregation

This is illustrated as Region 1 (Aggr) in the figure above.

Type II: Execution-Validity Constraints (Hard)

These constraints are required for correct execution. Extensions must preserve them to ensure that plans remain runnable.

Example: EnsureRequirements

Consider a partitioned hash join.

Before enforcing requirements:

HashJoinExec(mode=Partitioned, on=[t1.k = t2.k])
  Scan(t1)  <-- Not hash-partitioned on t1.k
  Scan(t2)  <-- Not hash-partitioned on t2.k

The plan represents the intended computation, but its children do not satisfy the hash join's distribution requirements.

After EnsureRequirements:

HashJoinExec(mode=Partitioned, on=[t1.k = t2.k])
  RepartitionExec(hash(t1.k), N)
    Scan(t1)
  RepartitionExec(hash(t2.k), N)
    Scan(t2)

Now the join's distribution requirements are satisfied.

We can express this as an invariant region, as shown in the figure above.

Within this region, every optimizer rule must preserve the invariant. If a rule violates it, the checker reports the failure immediately after that rule (when validation is enabled).

For extensions, the responsibility is explicitly one of the following:

1. Preserve the existing ordering/distribution guarantees.
2. If a transformation breaks those guarantees, repair the plan internally.
3. Alternatively, explicitly leave the region, perform transformations without those guarantees, and restore them before reentering. (demonstrated in the figure above)

What changes are included in this PR?

1. Add a configuration option, validate_optimizer_invariants, to control invariant checking. Use a feature flag to enable validation in tests/CI while skipping it in release builds.
2. The main change is in physical-optimizer/src/optimizer.rs, introducing invariant regions into the existing optimizer rule list.

The PoC implements the following regions, which are not exactly the same as the examples above:

• NoEmbeddedProjection: Soft constraint for projection shape canonicalization.
• HashJoinModeResolved: Soft constraint for hash join mode resolution.
• InputRequirementsSatisfied: Hard constraint; see the example above.

One thing that remains implicit is which optimizer rules establish the invariants.

For example, in the figure above, the canonical aggregation region is established by initial physical planning, while the InputRequirementsSatisfied region is established by EnsureRequirements.

Therefore, downstream custom optimizer pipelines must follow the same assumptions and manually declare these regions after the corresponding rules have established their invariants.

What is the testing strategy for this PR?

Are there any user-facing changes?

[alamb]

Thank you so much for working on this @2010YOUY01

The default physical optimizer pipeline works for the optimizations implemented in DataFusion core. However, there are many hidden assumptions between optimizer rules that make the pipeline work as a whole (e.g., rule A must run before rule B to produce a valid plan).

I think I fundamentally disagree with this assesment

what I don't like about this PR is that I think it makes the optimizer more complicated and commits us to non trivial invariants / rules going forward rather than driving towards something simpler.

Rather than encoding the set of existing (and I admit complicated) interconnections, I think the optimizer will be far simpler and easier to understand if there is one rule (in theory) to make the plan valid: EnforceRequirements

Then every other rule would obey "valid plan in --> valid plan out", which is a simple contract that is easy to explain and will make the Optimizer more understandable. I understand that the current code does not follow this ideal, but I feel like adding the API in this PR moves us away from a simpler future

In terms of this PR, I think that roughly boils down to always check these

                   Step::Begin(Arc::new(InputRequirementsSatisfied)),
                    Step::End(InputRequirementsSatisfied.name()),

In my opinion, the other examples in this PR such as

            Step::Begin(Arc::new(HashJoinModesResolved)),
            Step::Begin(Arc::new(NoEmbeddedProjections)),

should be treated like bug / limitation on rules and that we should work to remove these assumptions rather than making them explicit in the pipelines / assumptions

[2010YOUY01]

Thanks for sharing your perspective, @alamb!

I agree we should be more selective about adding invariant regions. EnsureRequirements is an obviously good candidate.

what I don't like about this PR is that I think it makes the optimizer more complicated and commits us to non trivial invariants / rules going forward rather than driving towards something simpler.

I'm not convinced that more explicit invariants necessarily mean more complexity. However, since I have less experience working with optimizers, I'll trust your judgment here and look for more problem-specific, compelling invariant regions to add in the future.

My current intuition is that more explicit constraints can actually reduce uncertainty and make extending the optimizer easier:

• Without invariant regions: When adding a new rule, it's hard to know where to insert it safely, since the existing assumptions between rules are mostly implicit or, at best, documented in comments.
• With invariant regions: The explicit constraints narrow down the possible insertion points to fewer, safer candidate slots, and automatic checks can catch violations of those assumptions.

When we find some invariants getting in the way, we can remove them relatively cheaply. Making them explicit also encourages reviewers to think more carefully about the design and existing assumptions.

Next step

I'll continue the discussion in #25688.

Generally, I think it's better to enforce the EnsureRequirements invariant rather than just document it as part of an analyzer/optimizer split.

However, enforcing it seems to require machinery of similar complexity to this PR. If we're going to introduce that machinery, perhaps we should make it extensible to other invariants as well, and see if we can add a few major, less controversial invariant fences in the future.

[alamb]

I'm not convinced that more explicit invariants necessarily mean more complexity. However, since I have less experience working with optimizers, I'll trust your judgment here and look for more problem-specific, compelling invariant regions to add in the future.

To be clear, I don't think the problem is with explicit invariants -- the problem is with some of the invariants themselves (e.g. the assumption that there are no embedded projections). I would rather remove as many of the invariants as possible rather than try and make them explicit

• Without invariant regions: When adding a new rule, it's hard to know where to insert it safely, since the existing assumptions between rules are mostly implicit or, at best, documented in comments.

Yes, this is a real problem -- but I think a better solution is to remove the implicit assumptions rather make them explicit

Generally, I think it's better to enforce the EnsureRequirements invariant rather than just document it as part of an analyzer/optimizer split.

I think it is basically enforced via https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html#method.check_invariants / https://docs.rs/datafusion-physical-optimizer/55.2.0/datafusion_physical_optimizer/sanity_checker/struct.SanityCheckPlan.html

As well as each pass here

datafusion/datafusion/core/src/physical_planner.rs

Lines 3377 to 3379 in b6f6990

	// This only checks the schema in release build, and performs additional checks in debug mode.
	OptimizationInvariantChecker::new(optimizer)
	.check(&new_plan, &before_schema)?;