SubhamSinghal opened a new pull request, #26417: URL: https://github.com/apache/datafusion/pull/26417
## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/26416 - Part of #24608. ## Rationale for this change With `datafusion.optimizer.enable_window_topn = true`, some top-K window queries return wrong results: SELECT pk, id, rn, rk FROM ( SELECT pk, v, id, row_number() OVER (PARTITION BY pk ORDER BY v) AS rn, rank() OVER (PARTITION BY pk ORDER BY v) AS rk FROM t) WHERE rk <= 1; returns only one row per partition and drops the rows tied at `rk = 1`. With the two functions swapped and `WHERE rn <= 1`, it returns rows with `rn = 2`. `WindowTopN` reads the filter's column index from the predicate, which indexes the projection between the filter and the window, and used it as if it indexed the window's output. When that projection drops a column (here `v`, which the outer query doesn't select), every later index shifts down by one. With two ranking functions the rule then builds `PartitionedTopKExec` from the wrong one, and since it also removes the `FilterExec`, nothing corrects the result. With a single ranking function, the same shifted index made the rule decline instead. So queries such as `row_number() OVER (PARTITION BY pk ORDER BY v, s)` with `s` unused afterwards never got the optimization. ## What changes are included in this PR? - New `column_index_below`, which maps the filter's column index down through the nodes `find_window_below` walks: a `ProjectionExec` must pass the column through as a bare `Column`, and the index follows it; a `RepartitionExec` leaves it unchanged; a column the projection computes (e.g. `rk + 1`) makes the rule decline. - `try_transform` resolves the index before checking that it names a window column. Everything after that already reads the resolved window expression. - The rule's doc comment lists the new condition. ## What is the testing strategy for this PR? `window_topn.slt`, each case checking both the rows (identical to `enable_window_topn = false`) and the `fn=` in `EXPLAIN`: - the two wrong-result queries above (`fn=rank`, `fn=row_number`), and a `dense_rank()` variant; - `row_number() OVER (PARTITION BY pk ORDER BY v, id)` with `v` dropped by the outer query, which is now rewritten; - `rank() OVER (...) + 1 AS x ... WHERE x <= 2`, which is still not rewritten: filter pushdown moves the expression into the predicate. `core/tests/physical_optimizer/window_topn.rs`, with hand-built plans for shapes the default optimizer pipeline doesn't produce but the rule accepts: - two stacked projections, which only resolve correctly top-down; - a projection over a `RepartitionExec` (the rule runs before `EnsureRequirements` inserts any); - a column the projection computes, e.g. `rk + 1`: unchanged plan. ## Are there any user-facing changes? With `enable_window_topn = true`, the affected queries now return correct results, and more query shapes are rewritten to `PartitionedTopKExec`. That includes `dense_rank()`, whose memory use is tracked in #. No API changes. The option is off by default. -- 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]
