[EPIC] Memory-efficient dedup on wide payload: fix ROW_NUMBER()=1 and DISTINCT ON patterns
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
## Motivation
Two real-world scenarios hit unnecessary OOM / severe memory blowup for the same class of query — "keep one row per group, pick winner by an ordering column":
**Scenario A** — `SELECT DISTINCT ON (id) * ORDER BY ts LIMIT 10` (already reported as #16620)
- Peak memory 1.4 GB for ~500K rows × 14 columns
- ~4× slower than `SELECT DISTINCT *` on the same data
**Scenario B** — production batch pipeline running `SELECT * FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY p ORDER BY o) rn FROM t) WHERE rn = 1` over nested JSON with wide `List>` payload
- 11M rows after UNNEST, ~2.5 KB per row
- OOMs at ~90 GB peak on 16 CPU (dominated by `ExternalSorterMerge`)
- Manual SQL rewrite into a two-step MAX + JOIN + FIRST_VALUE-without-ORDER-BY variant reduces peak to **3.6 GB** and wall time to ~40s
- The naive rewrite `FIRST_VALUE(wide_col ORDER BY o) GROUP BY p` **blows up to ~132 GB** — measured by an OS memory kill — due to per-group `Accumulator` fallback on nested types
## Root cause
`first_value` / `last_value` with `ORDER BY` falls back to the per-group `Accumulator` path for nested types (`List`, `Struct`, `Map`, ...) because `groups_accumulator_supported()` in [`first_last.rs`](https://github.com/apache/datafusion/blob/main/datafusion/functions-aggregate/src/first_last.rs) whitelists only scalar primitives (Int, Float, Decimal, Timestamp, Utf8, Binary).
Per-group `Accumulator` stores state as `Vec`. For `List>` a `ScalarValue::List` is a heap-allocated deep clone of the list value. `update_batch` clones on every candidate row and compares. Cost scales as `#rows × #groups × sizeof(wide_ScalarValue)`, catastrophic on wide payload.
The columnar `GroupsAccumulator` path — one `ArrayRef` per expression, updated via `arrow::compute::take` in batches — would scale as `#groups × wide_row × #expressions`. For Scenario B: **~125 MB instead of 132 GB**.
Additionally, `ROW_NUMBER() OVER (...) WHERE rn = 1` has no logical rewrite today. Currently the only physical optimization is `PartitionedTopKExec` (via the opt-in `enable_window_topn`), but that operator has its own `#groups × K × wide_row` scaling problem — for Scenario B we measured OOM against a 16 GB pool even with `K = 1`.
## Sub-issues
**Part 1 — Extend `first_value` / `last_value` `GroupsAccumulator` to nested types** ([#23601](https://github.com/apache/datafusion/issues/23601), blocker for the rest)
- Support `List`, `LargeList`, `ListView`, `Struct`, `Map` in `groups_accumulator_supported()` for `first_value` / `last_value` with `ORDER BY`
- Implement columnar accumulator state: `ArrayRef` per expression, updated via `arrow::compute::take` when a candidate wins per group
- Tests covering `List`, `Struct`, nested `List>`
- **Immediately fixes #16620** and removes the memory blowup on wide payload for Scenario B — no SQL rewrite required
**Part 2 — Coalesce peer `FIRST_VALUE(...) ORDER BY ` expressions into a single struct accumulator** ([#23602](https://github.com/apache/datafusion/issues/23602), optional, but recommended)
- Detect `SELECT ..., FIRST_VALUE(a ORDER BY o), FIRST_VALUE(b ORDER BY o), ... GROUP BY p` where all peer `FIRST_VALUE` expressions share the same `ORDER BY`
- Rewrite as `FIRST_VALUE(NAMED_STRUCT(a, b, ...) ORDER BY o) GROUP BY p` internally
- Cuts N argmax scans to 1 pass over the input; particularly beneficial for wide payload
**Part 3 — Logical rewrite `Filter(row_number() = 1) → Aggregate(FIRST_VALUE(... ORDER BY))`** ([#23603](https://github.com/apache/datafusion/issues/23603), user-visible auto-optimization)
- Match `Filter(Column(rn) = Literal(1)) → Projection? → WindowAggr(row_number() PARTITION BY p ORDER BY o)`
- Also cover `rn <= 1` and `rn < 2` as equivalents
- Preconditions: single `ROW_NUMBER` window function; filter is `Column op Literal(1)`; no other consumers of `rn`
- Rewrite to `Aggregate(GROUP BY p, aggr=[first_value(cols ORDER BY o)])`
- Depends on Part 1 (otherwise emits slow `Accumulator` path); works best with Part 2
- Semantics: ordering-key ties remain "unspecified", matching existing `ROW_NUMBER` behavior
## Dependencies
```
Part 1 (GroupsAccumulator nested) — standalone, ships #16620 fix on its own
│
├── Part 3 (rewrite rule) — auto-fires for users who wrote ROW_NUMBER()=1
│
└── Part 2 (coalesce peer FIRST_VALUE) — optional but recommended
```
## Success criteria
| Case | Before | Target |
|------|--------|--------|
| #16620 `DISTINCT ON` (~500K rows × 14 cols) | peak 1.4 GB, ~2 s | peak < 100 MB, sub-second |
| Wide-payload dedup (11M rows × ~2.5 KB) | OOM ~90 GB | peak < 500 MB, no SQL change |
## Related
- #16620 — Performance of `distinct on (columns)` — this epic subsumes it
- #12252 — `max_by` aggregate — related aggregate primitive (merged)
- #23263 — Late materialization when LIMIT prunes heavily — different family (scan-level, requires source-side random access + column late decode); complementary but not blocking
- #21479 — `PartitionedTopKExec` (WindowTopN) — incompatible with wide payload on high group cardinality; this epic complements it
Contributor guide
Assessment
This issue has not been assessed yet.