apache / apache/datafusion

Coalesce peer FIRST_VALUE / LAST_VALUE ORDER BY expressions into a single struct accumulator

Open
#23,602 3 comments 0 reactions 1 assignee Claimed by @tohuya6 View on GitHub
enhancement performance
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?

When a query has multiple `FIRST_VALUE(... ORDER BY )` (or `LAST_VALUE`) expressions that all share the same `ORDER BY` key, DataFusion currently creates one independent accumulator per expression. Each independently:

- Scans every input row
- Compares against its own stored best ``
- Maintains its own per-group state

For a typical "top-1 per key" pattern with N wide columns, this means N independent argmax passes, each holding its own `#groups × wide_row` state and doing its own row-by-row comparisons.

Example — user writes:

```sql
SELECT
p,
FIRST_VALUE(a ORDER BY o DESC),
FIRST_VALUE(b ORDER BY o DESC),
FIRST_VALUE(c ORDER BY o DESC),
FIRST_VALUE(d ORDER BY o DESC)
FROM t
GROUP BY p
```

Even after the nested-type `GroupsAccumulator` fix lands, this still runs 4 parallel argmax passes and holds 4× the state, purely because the accumulators cannot see they share an `ORDER BY`.

## Describe the solution you'd like

Add a logical optimizer rule that coalesces peer `FIRST_VALUE` / `LAST_VALUE` expressions with an identical `ORDER BY` key into a single struct-valued accumulator:

**Input plan:**

```
Aggregate:
group_by=[p]
aggr=[first_value(a ORDER BY o DESC),
first_value(b ORDER BY o DESC),
first_value(c ORDER BY o DESC),
first_value(d ORDER BY o DESC)]
```

**Rewritten plan:**

```
Projection: p, agg.wrapped[a], agg.wrapped[b], agg.wrapped[c], agg.wrapped[d]
Aggregate:
group_by=[p]
aggr=[first_value(NAMED_STRUCT('a', a, 'b', b, 'c', c, 'd', d) ORDER BY o DESC) AS wrapped]
```

The projection unpacks the struct back to individual columns so the rewrite is transparent at the schema level.

Preconditions:

- All coalesced expressions share the same `ORDER BY` (same expressions, directions, nulls-first/last)
- All are `FIRST_VALUE` or all are `LAST_VALUE` (do not mix — different orderings under nulls, and mixing violates the shared-ORDER-BY assumption)
- None have `DISTINCT` or `FILTER` clauses (would break the shared scan)

Benefits:

- One argmax pass over the input instead of N
- One per-group state slot instead of N
- On wide payloads this is a significant memory reduction (the wide row is stored once per group, not N times)

## Describe alternatives you've considered

- **Leave as-is and rely on the nested-type `GroupsAccumulator` fix alone.** Correct memory bound, but still N× the CPU work. For wide-payload aggregates the extra passes are the dominant cost.
- **Push the coalescing into a physical rule.** More flexible but harder to reason about — logical rule composes better with rewrite rules downstream (e.g. the `ROW_NUMBER = 1` rewrite ideally emits already-coalesced form).

## Additional context

Companion epic: https://github.com/apache/datafusion/issues/23600

This rule is most impactful when combined with:
- Nested-type `GroupsAccumulator` support (blocker for wide types)
- The `ROW_NUMBER() = 1` → aggregate rewrite (would emit the pre-coalesced form directly)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.