feat(physical-optimizer): support FilterExec with embedded projection in WindowTopN

[zhuqi-lucas]

Which issue does this PR close?

• Follow-up to Perf: Window topn optimisation #21479 (WindowTopN optimizer rule).

Rationale for this change

WindowTopN::try_transform currently bails out unconditionally when the top FilterExec carries an embedded projection (filter.projection().is_some()):

// Don't handle filters with projections
if filter.projection().is_some() {
    return None;
}

In practice, DataFusion's filter/projection pushdown pass often collapses a downstream ProjectionExec INTO the FilterExec's projection field, so real plans that match the FilterExec → BoundedWindowAggExec → SortExec pattern in every other respect are silently skipped and fall back to a full sort.

Observed at our prod:

SELECT ..., ROW_NUMBER() OVER (PARTITION BY g, d ORDER BY t DESC) AS rn
FROM ...
WHERE rn = 1

matches this exact pattern, but the FilterExec ends up with a projection pushed in, so WindowTopN skips and the sort-based path runs. This defeats the point of #21479 for a common shape.

What changes are included in this PR?

• Capture the FilterExec's projection indices at the start of try_transform.
• Run the existing PartitionedTopKExec rewrite as usual.
• Re-apply the captured projection as an outer ProjectionExec at the end so the transformed plan preserves the original output schema.

Plan shape before (for a filter with projection = [0, 1] that drops the ROW_NUMBER column):

FilterExec(predicate=rn <= 3, projection=[0, 1])
  BoundedWindowAggExec(ROW_NUMBER PBY pk OBY val)
    SortExec(pk, val)

Plan shape after:

ProjectionExec(expr=[pk@0, val@1])
  BoundedWindowAggExec(ROW_NUMBER PBY pk OBY val)
    PartitionedTopKExec(fetch=3, partition=[pk], order=[val], fn=row_number)

Are these changes tested?

Yes. Added a regression test filter_with_projection_still_rewrites in datafusion/core/tests/physical_optimizer/window_topn.rs:

• Builds FilterExec(rn <= 3, projection=[0, 1]) → BoundedWindowAggExec → SortExec.
• Runs WindowTopN with enable_window_topn = true.
• Asserts (via insta::assert_snapshot!) the resulting plan is ProjectionExec → BoundedWindowAggExec → PartitionedTopKExec → PlaceholderRowExec.

Existing 13 tests continue to pass:

cargo test -p datafusion --test core_integration physical_optimizer::window_topn -- --nocapture
test result: ok. 14 passed; 0 failed; 0 ignored; 0 measured

Also verified cargo fmt --all and cargo build -p datafusion-physical-optimizer succeed.

Are there any user-facing changes?

Yes. Queries matching ROW_NUMBER() OVER (PARTITION BY ...) ... WHERE rn OP K where an earlier optimizer pass has embedded a projection into the FilterExec will now be rewritten to PartitionedTopKExec (previously they silently fell back to a full sort). Only enabled when datafusion.optimizer.enable_window_topn = true.

[Copilot]

Pull request overview

Enhances the WindowTopN physical optimizer rule to rewrite eligible plans even when FilterExec contains an embedded projection (introduced by earlier filter/projection pushdown), preserving the original output schema by re-applying that projection as an outer ProjectionExec.

Changes:

• Capture FilterExec embedded projection indices during WindowTopN::try_transform and re-apply them as an outer ProjectionExec after the rewrite.
• Generalize extract_window_limit’s predicate parameter type by importing PhysicalExpr.
• Add a regression test covering FilterExec(predicate, projection=[..]) → BoundedWindowAggExec → SortExec rewrites.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
datafusion/physical-optimizer/src/window_topn.rs	Preserve FilterExec embedded projections by wrapping the rewritten plan in a ProjectionExec.
datafusion/core/tests/physical_optimizer/window_topn.rs	Add regression test ensuring embedded projections don’t prevent the PartitionedTopKExec rewrite.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[kosiew]

@zhuqi-lucas,

Thanks for working on this.
Supporting FilterExec with embedded projections is a nice improvement, and the regression test covers the plan shape well.

I found one correctness issue that should be addressed before merging. The rewrite currently drops FilterExec::fetch, which can change query results.

[kosiew]

Nice improvement capturing and restoring the embedded projection.

One thing I noticed is that the rewrite removes the FilterExec but does not preserve its fetch. FilterExecBuilder supports both an embedded projection and with_fetch, and FilterExec::execute applies the fetch after evaluating the predicate.

For example, if a matching projected filter has fetch=1, this rewrite currently produces only the outer ProjectionExec over the rewritten window. That returns all rn <= K rows instead of just one.

Could we either preserve the fetch with an equivalent outer limit/fetch operator, or skip this rewrite when filter.fetch().is_some()?

It would also be great to add a regression test covering the projection plus fetch case that executes the plan, or otherwise verifies that the row limit is preserved.

[zhuqi-lucas]

Good catch — fixed in f6bd9dd, and it turned out to be wider than the projection path.

I went with the second option: try_transform now declines when filter.fetch().is_some(). Preserving the fetch with an outer limit is not a straight substitution — PartitionedTopKExec bounds rows per partition, while FilterExec::fetch is a limit over the filtered output, so reproducing it would mean adding a real limit operator and reasoning about where it sits relative to the window. Declining keeps the rewrite honest and loses only the narrow intersection of "embedded projection and fetch".

Worth flagging: this was not introduced here. try_transform on main never reads fetch either — it just happens to bail on projection().is_some() first, so a FilterExec with a fetch and no projection was already being rewritten with its fetch dropped. The guard sits at the entry and is unconditional, so it covers that case too.

Two regression tests, both asserting the plan comes back unchanged: filter_with_projection_and_fetch_is_declined and filter_with_fetch_and_no_projection_is_declined (the second is the pre-existing case).

Also rebased onto main and resolved the conflicts — find_window_below now returns a generic intermediates list rather than a single proj_between, so the rewrite composes with that. One existing snapshot needed updating: the SortExec under the window now stays in place, which the old expectation predated.

[github-actions]

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.

[zhuqi-lucas]

@zhuqi-lucas,

Thanks for working on this. Supporting FilterExec with embedded projections is a nice improvement, and the regression test covers the plan shape well.

I found one correctness issue that should be addressed before merging. The rewrite currently drops FilterExec::fetch, which can change query results.

Thank you @kosiew for review, addressed review comments and added more tests now.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 91.66667% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.66%. Comparing base (c3ef346) to head (d75a081).

Files with missing lines	Patch %	Lines
datafusion/physical-optimizer/src/window_topn.rs	91.66%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23599      +/-   ##
==========================================
- Coverage   82.66%   82.66%   -0.01%     
==========================================
  Files        1147     1147              
  Lines      446357   446377      +20     
  Branches   446357   446377      +20     
==========================================
+ Hits       368971   368987      +16     
- Misses      54997    54998       +1     
- Partials    22389    22392       +3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.