Logical rewrite: Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY))
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
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
Contributor guide
Assessment
This issue has not been assessed yet.