david-mollitor-db opened a new pull request, #58864:
URL: https://github.com/apache/spark/pull/58864
### What changes were proposed in this pull request?
`TakeOrderedAndProjectExec` implements top-N (`ORDER BY ... LIMIT n`). For
unsorted input it selects
the top-K with `Utils.takeOrdered(iter.map(_.copy()), limit)(ord)`. The
`iter.map(_.copy())` copies
**every** input row, even though only `limit` rows are ever retained. The
copy is forced by the
selector: `org.apache.spark.util.collection.Utils.takeOrdered` wraps Guava's
`Ordering.leastOf` /
`TopKSelector`, which buffers element *references* while selecting, and the
child's
whole-stage-codegen iterator yields a single reused `UnsafeRow` -- so every
row must be copied to a
detached instance before being offered. For N input rows we allocate N
copies to keep `limit`.
This replaces the copy-everything-then-select approach with a bounded top-K
that copies a row only
when it is actually retained. A helper on the `TakeOrderedAndProjectExec`
companion builds a bounded
max-heap (`java.util.PriorityQueue` with `ord.reverse`, so the head is the
current largest of the
retained set = the eviction threshold), compares each *live* row against the
head, and calls
`copy()` only on insertion:
```scala
if (heap.size < num) heap.add(row.copy())
else if (ord.compare(row, heap.peek()) < 0) { heap.poll();
heap.add(row.copy()) }
```
Rows are drained into ascending order to match the sorted output of the
previous
`Utils.takeOrdered`. This is applied to the three copy-all sites (the
`executeCollect` unsorted
branch, the `doExecute` per-partition local top-K, and the `doExecute`
post-shuffle merge). The
already-efficient `orderingSatisfies` fast paths
(`_.map(_.copy()).take(limit)`) and the post-limit
OFFSET/projection are unchanged. Copies drop from O(N) to O(limit); a
per-element heap decision
against the exact current threshold enables copy-on-retain, which Guava's
reference-buffering design
cannot.
### Why are the changes needed?
`ORDER BY ... LIMIT n` over large unsorted input is common (top-N, "latest
100", dashboards). JFR
profiling of `spark.range(20000000).selectExpr("id","id % 1000 as
k").orderBy("k").limit(100)` showed
`UnsafeRow.copy()` at ~93.7% of sampled allocation and ~15.5% of CPU -- the
operator's cost is almost
entirely copying rows that are immediately discarded. Cutting copies from
O(N) to O(limit) removes
that.
Before/after on that query (local run):
| Metric | Before (copy-all) | After (copy-on-retain) |
|---|---|---|
| `UnsafeRow.copy()` allocation samples | 2532 / 2829 | 4 / 2342 |
| `orderBy(k).limit(100)`, noop | 671 ms | 343 ms |
| `orderBy(k).limit(100)`, collect | 650 ms | 376 ms |
The allocation drop is structural; the ~2x wall-time is a local measurement
(fewer copies = less
allocation, GC, and memory bandwidth).
### Does this PR introduce _any_ user-facing change?
No. Both old and new return the `limit` smallest rows by the ordering, in
ascending order; which
rows win a tie at the exact boundary value is unspecified in both (inherent
LIMIT non-determinism).
`limit`/OFFSET semantics are identical (the same `limit` is passed and the
post-limit `.drop(offset)`
is untouched). Worst-case memory is unchanged: `limit` is bounded by
`spark.sql.execution.topKSortFallbackThreshold`, and the heap holds at most
`limit` rows (vs the
previous selector's ~2*`limit` buffered references).
### How was this patch tested?
`TakeOrderedAndProjectSuite` (with/without projection, 0/1/10 partitions, 0
and 10k rows, the
already-sorted path via both `executeCollect` and `doExecute`),
`PlannerSuite` (including the
`topKSortFallbackThreshold` and `limit + offset` threshold planner tests),
and
`InsertSortForLimitAndOffsetSuite` all pass unchanged.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.
--
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]
