zhuqi-lucas opened a new pull request, #23599:
URL: https://github.com/apache/datafusion/pull/23599
# Which issue does this PR close?
- Follow-up to #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()`):
```rust
// 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 Massive (atlas / polygon): the production MV shape
```sql
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`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
