kosiew commented on code in PR #23869:
URL: https://github.com/apache/datafusion/pull/23869#discussion_r3681020879
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1811,6 +1811,394 @@ impl PartitionedTopKRank {
}
}
+/// Per-partition state for `DENSE_RANK()` semantics.
+///
+/// A `HashMap<Vec<u8>, Vec<TieEntry>>` keyed by the row-encoded ORDER BY
+/// bytes, capped at `k` distinct keys. Each key's `Vec<TieEntry>` holds
+/// every row seen at that ob value, one entry per contributing source
+/// `RecordBatch`. `TieEntry` is reused verbatim from [`RankPartitionState`].
+struct DenseRankPartitionState {
+ groups: HashMap<Vec<u8>, Vec<TieEntry>>,
+ /// Cached admission boundary: the largest tracked ob value once
+ /// `groups` is full. A `HashMap` is unordered, so finding it otherwise
+ /// costs an O(K) scan of the keys on *every* new distinct ob value —
+ /// O(N·K) overall on mostly-distinct input. Caching makes the common
+ /// "new ob is worse than the boundary → drop" path O(1); the O(K) scan
+ /// is paid only when the cache is stale (`None`) — after a fill or an
+ /// eviction changes the key set.
+ max_key: Option<Vec<u8>>,
+}
+
+impl DenseRankPartitionState {
+ fn size(&self) -> usize {
+ let table_overhead =
+ self.groups.capacity() * (size_of::<Vec<u8>>() +
size_of::<Vec<TieEntry>>());
+ let contents: usize = self
+ .groups
+ .iter()
+ .map(|(key, entries)| {
+ key.capacity()
+ + entries.capacity() * size_of::<TieEntry>()
+ + entries
+ .iter()
+ .map(|e| {
+ e.row_indices.capacity() * size_of::<u32>() +
e.batch_bytes
Review 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.
```bash
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.
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1811,6 +1811,394 @@ impl PartitionedTopKRank {
}
}
+/// Per-partition state for `DENSE_RANK()` semantics.
+///
+/// A `HashMap<Vec<u8>, Vec<TieEntry>>` keyed by the row-encoded ORDER BY
Review 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.
--
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]
