Logical rewrite: Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY))

[zhuqi-lucas]

Part of #23600.

Is your feature request related to a problem or challenge?

A very common "top-1 per group" pattern is written as:

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), 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: #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:

• Perf: Window topn optimisation #21479 — PartitionedTopKExec — a complementary physical rewrite; the aggregate rewrite is the natural fit for wide payloads / high group cardinality

[saadtajwar]

take