SubhamSinghal opened a new issue, #26876:
URL: https://github.com/apache/datafusion/issues/26876
### Describe the bug
With `datafusion.optimizer.enable_window_topn = true`, a `DENSE_RANK()`
top-K query holds memory proportional to its **input**, not to the rows it
retains, for every payload type (Int64 included).
`PartitionedTopKDenseRank::insert_batch` registers the whole input batch in
its `RecordBatchStore`
(`self.store.register(batch.clone())`), and the batch stays there as long as
any retained row references it. There
is no gather of the admitted rows and no compaction, so one surviving row
keeps its whole batch alive. `ROW_NUMBER` and `RANK` avoid this through
`PartitionedTopKCore`: they gather only the admitted rows and compact the store
on a row ratio (#25730, #25968). `DENSE_RANK` hasn't been moved over.
### To Reproduce
In `datafusion-cli`, create 4.1M rows in which each 8192-row batch
introduces 10 new partitions (same data as #26381):
COPY (
SELECT CAST((i / 8192) * 10 + (i % 10) AS BIGINT) AS pk,
CAST(i AS BIGINT) AS v,
concat(md5(CAST(i AS VARCHAR)), md5(CAST(i + 1 AS VARCHAR)),
md5(CAST(i + 2 AS VARCHAR)), substr(md5(CAST(i + 3 AS
VARCHAR)), 1, 18)) AS s
FROM generate_series(0, 4095999) AS t(i)
ORDER BY i
) TO '/tmp/ptk.parquet' STORED AS PARQUET;
CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION '/tmp/ptk.parquet';
SET datafusion.execution.target_partitions = 1;
SET datafusion.optimizer.enable_window_topn = true;
SELECT count(*), sum(length(s)), sum(pk), sum(v), sum(rn)
FROM (SELECT pk, v, s, dense_rank() OVER (PARTITION BY pk ORDER BY v) AS
rn FROM t)
WHERE rn <= 1;
The query retains 5,000 rows. Smallest passing `--memory-limit`:
| | memory limit |
|---|---|
| `dense_rank()`, `enable_window_topn = true` | **1046M** |
| `dense_rank()`, `enable_window_topn = false` (sort plan, spills) | 17M |
At the operator level, with K=1 and the same 200 batches of 8192 rows (10
survivors spread through each batch)
fed to all three operators, the stores hold:
| payload | ROW_NUMBER | RANK | DENSE_RANK | input |
|---|---|---|---|---|
| Int64 | 31 KB (2,002 rows) | 31 KB (2,001 rows) | **25,600 KB (1,638,400
rows)** | 25,600 KB |
| Utf8 | 247 KB | 246 KB | **224,012 KB** | 224,012 KB |
### Expected behavior
Retained memory proportional to the retained rows (the rows of K distinct
ORDER BY values per partition), as with `ROW_NUMBER` and `RANK`, and no more
than the sort plan the rewrite replaces.
### Additional context
Gathering the admitted rows alone isn't enough. With survivors spread across
batches the store still grows with
the input unless it is also compacted; that is why `PartitionedTopKCore`
compacts at `STORE_COMPACTION_RATIO`. One difference from `RANK`: `DENSE_RANK`
keeps every row of each of its K retained ORDER BY values, so an evicted value
releases all of its rows at once.
--
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]
