SubhamSinghal opened a new pull request, #25730:
URL: https://github.com/apache/datafusion/pull/25730
## Which issue does this PR close?
- Part of #24608. Close #13586.
## Rationale for this change
`enable_window_topn` rewrites `Filter(rn<=K) → Window(ROW_NUMBER) → Sort`
into
`PartitionedTopKExec`. As its own config documentation warns, it helps at
low partition cardinality
and regresses at high cardinality. Running the existing ROW_NUMBER sweep in
`benchmarks/queries/h2o/window.sql` on `J1_1e7_1e7_NA` (10M rows) and
comparing the flag on vs off
on the same commit (medians of 7):
| partitions | CPU, ON vs OFF on `main` |
|---|---|
| 100 | 5.65× faster |
| 1,000 | 5.21× faster |
| 10,000 | 2.28× faster |
| 100,000 | **2.16× slower** |
The benefit decays with cardinality and inverts around 100K, so the
optimization becomes a
pessimization exactly where a per-partition top-K should help most. That is
why the flag still
defaults to `false`.
The cause is how retained rows are stored, not the per-partition heaps. For
every
(batch, partition) pair, `insert_batch` `take`s a sub-batch of *all* that
partition's rows in that
batch — qualifying or not — registers it in a `RecordBatchStore` private to
that partition, then
runs `maybe_compact()`. At 100K partitions that is ~100K separate stores,
one `take` per
(batch, partition), and a compaction check per (batch, partition) on every
batch.
## What changes are included in this PR?
Two coupled changes to the `ROW_NUMBER` arm of `PartitionedTopKExec`. `RANK`
and `DENSE_RANK` are
untouched — see the follow-ups below.
1. **One `RecordBatchStore` for the operator** instead of one per partition.
Each partition holds a
lean `PartitionHeap` of `PartitionSlot { key: Vec<u8>, batch_id: u32,
row: u32 }` — the
row-encoded ORDER BY key inline, output columns by reference, refcounted
as `DENSE_RANK` already
does.
2. **Decide first, gather once.** A single pass over the batch admits rows
and records only their
indices; then **one** `take_record_batch` of just the admitted rows
becomes one store entry. If
every admission from a batch is superseded within that batch, nothing is
registered at all.
Sharing the store means `TopKHeap::maybe_compact` can no longer be used on
this path — it calls
`store.clear()`, which on a shared store invalidates *other* partitions'
`batch_id`s, and its guard
reads a now-global `total_rows`. Since that per-heap compaction was the only
thing bounding retention
on `main`, it is replaced by an operator-level one:
3. **`compact_store()`**, run at the end of every `insert_batch`, rewrites
the store when it holds
more than `STORE_COMPACTION_RATIO` × the rows the heaps actually point
at. Without it, retention is
**unbounded**.
## What is the testing strategy for this PR?
`datafusion-physical-plan` (1964 tests) and all 512 sqllogictest files pass.
plus added new UTs
## Are there any user-facing changes?
No API changes and no change to query results. `enable_window_topn` still
defaults to `false`.
## Benchmark details
Queries are the ROW_NUMBER cardinality sweep already in
`benchmarks/queries/h2o/window.sql`
(k=2, `ORDER BY v2 DESC`, `WHERE v2 IS NOT NULL`), run against
`J1_1e7_1e7_NA.parquet`
| partitions | main, ON vs OFF | this branch, ON vs OFF |
|---|---|---|
| 100 | 5.65× faster | 5.80× faster |
| 1,000 | 5.21× faster | 5.90× faster |
| 10,000 | 2.28× faster | 4.59× faster |
| 100,000 | **2.16× slower** | **2.01× faster** |
**"Why not just extract `RecordBatchStore` out of `TopKHeap`?"** — measured,
on a throwaway branch
with a shared store and `main`'s per-(batch, partition) gather. Sharing the
store is ~1.2× of the
win and is *worse* on RSS than `main` (it pins per-(batch, partition)
sub-batches with nothing
reclaiming them); decide-then-gather is the remaining ~5.3×, and without it
the rewrite is still
slower than the `SortExec` plan it replaces.
--
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]
