zhuqi-lucas opened a new issue, #23600: URL: https://github.com/apache/datafusion/issues/23600
## 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<Struct<...>>` 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<ScalarValue>`. For `List<Struct<...>>` 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 **#1 — Extend `first_value` / `last_value` `GroupsAccumulator` to nested types** (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<Utf8>`, `Struct<int, utf8>`, nested `List<Struct<...>>` - **Immediately fixes #16620** and removes the memory blowup on wide payload for Scenario B — no SQL rewrite required **#2 — Coalesce peer `FIRST_VALUE(...) ORDER BY <o>` expressions into a single struct accumulator** (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 **#3 — Logical rewrite `Filter(row_number() = 1) → Aggregate(FIRST_VALUE(... ORDER BY))`** (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 #1 (otherwise emits slow `Accumulator` path); works best with #2 - Semantics: ordering-key ties remain "unspecified", matching existing `ROW_NUMBER` behavior ## Dependencies ``` #1 (GroupsAccumulator nested) — standalone, ships #16620 fix on its own │ ├── #3 (rewrite rule) — auto-fires for users who wrote ROW_NUMBER()=1 │ └── #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 -- 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]
