zhuqi-lucas opened a new issue, #23603: URL: https://github.com/apache/datafusion/issues/23603
Part of https://github.com/apache/datafusion/issues/23600. ## Is your feature request related to a problem or challenge? A very common "top-1 per group" pattern is written as: ```sql SELECT * FROM ( SELECT *, ROW_NUMBER() OVER (PARTITION BY p ORDER BY o DESC) AS rn FROM t ) WHERE rn = 1 ``` DataFusion currently plans this as: ``` FilterExec(rn = 1) BoundedWindowAggExec(row_number() PARTITION BY p ORDER BY o) SortExec(p, o DESC) ← buffers all input rows ... ``` The `SortExec` buffers the entire input, including all wide payload columns, even though the semantics only require one row per `p`. On wide payload this scales badly: - Production observation: 11M rows × ~2.5 KB payload → ~90 GB peak OOM on 16 CPU - The hand-written aggregate equivalent (`SELECT p, FIRST_VALUE(...ORDER BY o DESC), ... GROUP BY p`) runs at 3.6 GB peak on the same data, once the nested-type `GroupsAccumulator` blocker is out of the way The only physical-side optimization today is `PartitionedTopKExec` (`enable_window_topn`, [#21479](https://github.com/apache/datafusion/pull/21479)), which does not compose well with wide payloads on high-cardinality groups (each group carries K wide rows in a heap → still `#groups × K × wide_row`). ## Describe the solution you'd like Add a logical optimizer rule that rewrites the `ROW_NUMBER() = 1` pattern to an aggregate: **Input plan:** ``` Filter: rn = 1 -- or rn <= 1, or rn < 2 Projection: ... -- optional passthrough WindowAggr: row_number() OVER (PARTITION BY p ORDER BY o DESC) AS rn child ``` **Rewritten plan:** ``` Aggregate: group_by=[p] aggr=[first_value(col_i ORDER BY o DESC) for each output col_i] child ``` Preconditions for the rewrite: - The window function is exactly `row_number()` (not `rank`, `dense_rank`, `nth_value`, etc.) - Filter predicate is `rn = 1`, `rn <= 1`, or `rn < 2` (all equivalent to "top-1 per group") - `rn` is not referenced anywhere else in the plan above the filter (no aliasing, no join key, no other projection) - The window has a `PARTITION BY` clause (without it the semantics differ) Behavior: - Order-key ties remain "unspecified", matching the existing `ROW_NUMBER` behavior (both `ROW_NUMBER` and `FIRST_VALUE` break ties arbitrarily) - All non-partition columns are wrapped in `FIRST_VALUE(... ORDER BY o)` in the aggregate output ## Describe alternatives you've considered - **Physical rule instead of logical.** Feasible but constrained — a logical rule composes with downstream rules (projection pushdown, partial aggregate pushdown, distribution planning) that don't know about the physical `BoundedWindowAggExec` shape. - **Expand `PartitionedTopKExec` to be memory-efficient on wide payloads.** Would need late materialization or row-id primitives that DataFusion does not have today. The aggregate rewrite reuses existing infrastructure. - **Leave to users.** Users routinely write the `ROW_NUMBER() = 1` form (it is the standard SQL idiom for "pick one row per group with ordering"); auto-rewrite converts a well-known trap into transparent optimization. ## Additional context Companion epic: https://github.com/apache/datafusion/issues/23600 This rewrite depends on: - Nested-type `GroupsAccumulator` support — without it the rewrite would emit the slow per-group `Accumulator` path on wide payloads and could regress users. **Rewrite should only fire when the aggregate destination is confirmed to hit the columnar path**, or should be gated by config until the blocker lands. - Peer `FIRST_VALUE` coalescing — optional but strongly recommended for wide-payload rewrites (avoids N × per-group state). Related: - #21479 — `PartitionedTopKExec` — a complementary physical rewrite; the aggregate rewrite is the natural fit for wide payloads / high group cardinality -- 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]
