perf: Extend WindowTopN to dense_rank - #23869
Conversation
github-actions
Bot
added
optimizer
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23869 +/- ##
==========================================
+ Coverage 80.72% 80.87% +0.15%
==========================================
Files 1090 1099 +9
Lines 370384 375428 +5044
Branches 370384 375428 +5044
==========================================
+ Hits 298975 303628 +4653
- Misses 53590 53637 +47
- Partials 17819 18163 +344 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
comphead
changed the title
dense_rank
|
@kumarUjjawal @kosiew can you please review this PR? |
kumarUjjawal
left a comment
kumarUjjawal
left a comment
There was a problem hiding this comment.
Hi @SubhamSinghal I have left some comments please let me know what you think.
| match udf.fun().name() { | ||
| "row_number" => Some(WindowFnKind::RowNumber), | ||
| "rank" => Some(WindowFnKind::Rank), | ||
| "dense_rank" => Some(WindowFnKind::DenseRank), |
There was a problem hiding this comment.
Could we skip this rewrite when another expression in the same window operator needs following rows? The input is pruned before all window expressions run, so a sibling lead() returns the wrong value at the retained boundary. We can add a dense_rank with lead regression and only rewrite when every sibling expression is safe under tail pruning.
There was a problem hiding this comment.
Fixed in 4530bba. The rewrite now bails unless every window expr is ROW_NUMBER/RANK/DENSE_RANK over the same PARTITION BY/ORDER BY. LEAD and aggregates are excluded. Added EXPLAIN + correctness SLT regressions for SUM and LEAD siblings.
| // For RANK / DENSE_RANK: an empty ORDER BY makes every row tie | ||
| // at rank 1 — the optimization is degenerate (we'd retain the | ||
| // entire input) and tie storage would be unbounded. | ||
| if matches!(fn_kind, WindowFnKind::Rank | WindowFnKind::DenseRank) |
There was a problem hiding this comment.
This guard misses orderings that become empty after partition keys are removed. DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk) reaches PartitionedTopKExec without order expressions and panics. Please validate the effective order keys and add a regression.
There was a problem hiding this comment.
Fixed in 4530bba. Replaced order_by().is_empty() with a structural guard: sort_exec.expr().len() <= partition_prefix_len → bail. Catches PARTITION BY pk ORDER BY pk (and no-ORDER-BY ROW_NUMBER). Added SLT regression.
| entries.push(TieEntry { | ||
| batch: batch.clone(), | ||
| row_indices: run_indices, | ||
| batch_bytes: input_batch_bytes, |
There was a problem hiding this comment.
records the full source-batch size for every retained group even though the batch clones share Arrow buffers. With many partitions, memory reservation is heavily inflated and can fail with ResourcesExhausted. This expands #23326 to the normal dense-rank path; please account each source batch once or compact the retained rows.
There was a problem hiding this comment.
that will leaves the normal dense-rank path counting one shared batch once per retained group across every partition. This can greatly inflate the reservation and cause false ResourcesExhausted errors. Could we fix the shared-batch accounting here, or keep this rewrite disabled until #23326 is fixed?
| .groups | ||
| .keys() | ||
| .map(|key| key.as_slice()) | ||
| .max() |
There was a problem hiding this comment.
Finding the maximum scans all K keys for every new distinct value. On mostly distinct input this makes the rewrite O(N × K) and can be much slower for large K than the sort it replaces.
There was a problem hiding this comment.
I had kept it unoptimised thinking K will be small. Fixed in 4530bba. Cache the admission boundary (max_key); the O(K) scan runs only when stale (after eviction/fill), so the common "worse than boundary → drop" path is O(1).
| if is_smaller { | ||
| // Evict the entire max-key group and drop the cached | ||
| // max (recomputed on the next Case C). | ||
| let evicted_key = state.max_key.take().expect("max key present"); |
There was a problem hiding this comment.
The cache only makes rejected candidates O(1). Every accepted candidate invalidates it, so input arriving from worst to best still scans all K keys for each new value and remains O(N × K). The benchmarks use K=2 and do not exercise this case. Could we use an ordered structure or add a large-K improving-order benchmark?
There was a problem hiding this comment.
Fixed in a74ed66. Replaced with binaryHeap so that worst case will be 0(Nxlog(K)) instead 0(NxK)
| entries.push(TieEntry { | ||
| batch: batch.clone(), | ||
| row_indices: run_indices, | ||
| batch_bytes: input_batch_bytes, |
There was a problem hiding this comment.
that will leaves the normal dense-rank path counting one shared batch once per retained group across every partition. This can greatly inflate the reservation and cause false ResourcesExhausted errors. Could we fix the shared-batch accounting here, or keep this rewrite disabled until #23326 is fixed?
|
@2010YOUY01 Can you take a look at this whenever you get the chance. |
There was a problem hiding this comment.
@SubhamSinghal
Thanks for the dense-rank implementation. I think the overall approach looks good, but I found one memory accounting issue that can lead to very large overestimates and trigger unnecessary OOMs. I also have one small naming suggestion that could make the ownership and accounting model a bit clearer.
| + entries | ||
| .iter() | ||
| .map(|e| { | ||
| e.row_indices.capacity() * size_of::<u32>() + e.batch_bytes |
There was a problem hiding this comment.
I think there is a memory accounting issue here. DenseRankPartitionState::size() adds e.batch_bytes for every TieEntry, but the new dense-rank path creates one TieEntry per retained ORDER BY group while each entry only holds an Arc clone of the same source RecordBatch (see the insertions at lines 2056, 2068, and 2104).
That means a single input batch with K retained distinct values is charged roughly K * get_record_batch_memory_size(batch), even though the Arrow buffers are only retained once. The existing rank tie path has some precedent for conservative duplicate charging, but for dense-rank this becomes the common case because one batch can contribute many retained groups.
I reproduced this with datafusion-cli using the script below. With enable_window_topn=true, target_partitions=1, --batch-size 2000, and --memory-limit 50M, the query fails trying to allocate about 2001.6 MiB for PartitionedTopKDenseRank[0]. Running the same query with enable_window_topn=false succeeds and returns 200 rows.
python3 - <<'PY'
with open('/tmp/df_dr.csv','w') as f:
f.write('pk,val,payload\n')
payload='x'*10000
for g in range(1,1001):
f.write(f'{g%2},{g},{payload}\n')
PY
cargo run -q -p datafusion-cli -- \
--batch-size 2000 --memory-limit 50M \
--command "
SET datafusion.optimizer.enable_window_topn = true;
SET datafusion.execution.target_partitions = 1;
CREATE EXTERNAL TABLE dr(pk INT, val INT, payload VARCHAR)
STORED AS CSV LOCATION '/tmp/df_dr.csv'
OPTIONS (format.has_header true);
SELECT count(pk), count(val), count(payload)
FROM (
SELECT pk, val, payload,
DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) AS drnk
FROM dr
)
WHERE drnk <= 100;
"Could we account for retained batches uniquely instead? For example, by reusing or extending the existing RecordBatchStore style accounting, or by storing batch IDs with ref-counted byte accounting. It would also be great to add a regression test covering the case where one batch contributes multiple retained dense-rank groups.
There was a problem hiding this comment.
Fixed in 8e97623. Using RecordBatchStore as suggested.
|
|
||
| /// Per-partition state for `DENSE_RANK()` semantics. | ||
| /// | ||
| /// A `HashMap<Vec<u8>, Vec<TieEntry>>` keyed by the row-encoded ORDER BY |
There was a problem hiding this comment.
Small naming suggestion. The dense-rank implementation currently describes the retained groups as HashMap<Vec<u8>, Vec<TieEntry>>. After fixing the memory accounting, it might be worth introducing a type name that reflects retained batch slices instead of reusing TieEntry, since these entries represent more than just boundary ties. I think that would make the ownership and accounting invariant a bit clearer.
There was a problem hiding this comment.
Renamed it to GroupEntry
There was a problem hiding this comment.
@SubhamSinghal
Thanks for the follow-up. The heap-based boundary tracking looks promising, but I think there are still two memory-accounting issues to address before this is ready. The dense-rank state continues to charge shared RecordBatch bytes multiple times, and the new eager heap allocation can consume substantial memory when both k and partition cardinality are large.
| /// The same keys as `groups`, in a max-heap: the admission boundary | ||
| /// (the largest tracked ob value) is an O(1) `peek()` check, and | ||
| /// admission / removal are O(log K). | ||
| /// |
There was a problem hiding this comment.
I think this still over-counts retained memory. Each TieEntry only holds an Arc clone of a RecordBatch, but size() adds e.batch_bytes for every entry. In dense-rank processing, a single input batch can contribute several retained ORDER BY groups, so the same underlying batch bytes may be charged multiple times.
The insertion sites around lines 2062, 2075, and 2100 also continue to assign the full input_batch_bytes value to each entry. This can inflate the reservation and cause a false ResourcesExhausted error when enable_window_topn=true.
Could we account for each shared batch allocation only once, for example by tracking unique batch allocations or moving ownership of the batch-level charge outside individual tie entries?
There was a problem hiding this comment.
@kosiew are you referring to RANK code? at line 2100 I can see:
state.groups.insert(
ob_key,
vec![GroupEntry {
row_indices: run_indices,
batch_id,
}],
);
which is referring to batch_id. Looks like stale comment but please confirm.
| groups.clear(); | ||
| { | ||
| let pk_rows = &self.partition_scratch_rows; | ||
| for i in 0..num_rows { |
There was a problem hiding this comment.
Could we avoid allocating capacity for all k heap entries as soon as a partition is first seen? With a large k and many partitions, BinaryHeap::with_capacity(k) front-loads roughly O(partitions * k) heap-slot allocation even when most partitions retain only a few distinct keys.
Using BinaryHeap::new() or a smaller initial capacity would let the heap grow with actual usage and reduce reservation pressure for high-cardinality partition workloads.
There was a problem hiding this comment.
fixed in afc931c using BinaryHeap::new()
kumarUjjawal
left a comment
kumarUjjawal
left a comment
There was a problem hiding this comment.
Overall looks good, just one issue I found.
kumarUjjawal
left a comment
kumarUjjawal
left a comment
There was a problem hiding this comment.
Thank you @SubhamSinghal
4 participants
Which issue does this PR close?
ROW_NUMBER < 5/ TopK #6899. Completes the ROW_NUMBER + RANK + DENSE_RANK trio requested by the issue (ROW_NUMBER shipped in Perf: Window topn optimisation #21479; RANK is in flight as PR A).Rationale for this change
Naive plan for
Filter(dr <= K) → BoundedWindowAggExec(DENSE_RANK) → SortExec(pk, ob) → inputsorts the full input even though only rows at the K distinct-smallest ob values per partition are needed. DENSE_RANK is stable under tail-pruning, so the rewrite preserves ranks for every retained row.What changes are included in this PR?
Are these changes tested?
Yes
h2o bench — Q24–Q29 added; top-2 on 10M-row
large, 3-iter avg vsenable_window_topn=false:Are there any user-facing changes?
No breaking changes. Rule fires on DENSE_RANK when
datafusion.optimizer.enable_window_topn = true; flag staysfalseby default