SubhamSinghal opened a new pull request, #27136: URL: https://github.com/apache/datafusion/pull/27136
## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/26876 ## Rationale for this change With `datafusion.optimizer.enable_window_topn = true`, a `DENSE_RANK() ... WHERE rn <= K` query holds memory proportional to its **input**, not to the rows it keeps. `PartitionedTopKDenseRank` registers every whole input batch in its store and never compacts it, so one retained row keeps its whole batch alive. ROW_NUMBER and RANK avoid this through `PartitionedTopKCore` (#25730, #25968); DENSE_RANK was never ported. This blocks enabling the rule by default (#24608) independently of #26381. Smallest passing `--memory-limit`, `target_partitions = 1`, the #26381 data (4.1M rows, 5,000 retained): | payload | `dense_rank` before | `dense_rank` after | `row_number` | flag off (sort plan) | |---|---|---|---|---| | Utf8 | 542M | **4M** | 3M | 13M | | Int64 | 66M | **3M** | 2M | 165M | | Utf8View | 589M | 464M | 463M | 17M | What is left on Utf8View is #26381's view-buffer pinning, which `row_number` shows too. ## What changes are included in this PR? `PartitionedTopKDenseRank` now runs on `PartitionedTopKCore<DenseRankPartitionState>`, like the other two operators: - **One pass over the rows in batch order**, replacing the per-batch demux by partition and bucketing by ORDER BY value. After any prefix of a partition's rows the retained set is the rows of its K distinct-smallest values, so arrival order cannot change a decision. - **One store slot per retained row.** A group is a `Vec<StoreRef>`; the store gathers only the admitted rows and is compacted by `RecordBatchStore::compact`. Evicting a group releases each of its rows (rows from the batch in flight come off the pending count). - **`size()` is O(1).** A running `groups_bytes` that grows and shrinks, like RANK's `ties_bytes`, replaces the per-partition `contents_bytes` sum, which was O(partitions) per batch. The groups table is charged through `HashMap::allocation_size`: `capacity()` drops on removal (tombstones) without freeing anything, so a running total built on it would drift. - **Emit streams** through `PartitionedTopKCore::emit_with` / `EmitState`, instead of building every output batch before returning. - `row_replacements` keeps its DENSE_RANK meaning (rows evicted), but its values change: the old code visited a batch's ORDER BY values in hash order and often admitted then evicted a group within the same batch, which the row-order pass does not. On h2o-shaped input at 20% retention it drops from ~29K to ~4.5K with identical output. CPU, h2o `J1_1e7_1e7_NA.parquet`, the DENSE_RANK queries with the flag on, median of 7 paired reps (order alternating), before → after. `window` Q26/Q28 order by `v2 % 10`, which is float modulo on Float64 `v2` and leaves the values near-unique (both emit exactly K rows per partition, like Q25/Q27); the tie-heavy cases are `window_sorted` Q15/Q16 and the table below. | query | main | PR | speedup | |---|---|---|---| | window Q24, ~100 partitions | 29.3 | 29.5 | 0.99× | | window Q25, ~1K | 32.4 | 29.6 | 1.09× | | window Q26, ~1K, `v2 % 10` (no ties) | 40.9 | 36.9 | 1.11× | | window Q27, ~10K | 38.0 | 29.2 | 1.30× | | window Q28, ~10K, `v2 % 10` (no ties) | 46.2 | 36.2 | 1.28× | | window Q29, ~100K | 108.7 | 73.5 | 1.48× | | window_sorted Q11, 100 partitions | 24.3 | 23.9 | 1.02× | | window_sorted Q12, 1K | 34.1 | 21.3 | 1.60× | | window_sorted Q13, 10K | 46.5 | 23.3 | 2.00× | | window_sorted Q14, 100K | 59.5 | 27.3 | 2.18× | | window_sorted Q15, 1K heavy ties | 56.2 | 52.1 | 1.08× | | window_sorted Q16, 10K heavy ties | 73.0 | 62.1 | 1.18× | ## What is the testing strategy for this PR? - `test_partitioned_topk_dense_rank_bookkeeping_tracks_recompute` (replaces `..._contents_bytes_tracks_recompute`): 256 randomized shapes; after every batch it checks every slot against the store (`assert_store_matches_slots`), `groups_bytes` and `index_bytes` against a recompute, the reservation against `size()`, and `size()` against the core's terms plus the recomputed groups; then the output against brute force. It asserts the shapes both evicted and compacted. - `test_partitioned_topk_dense_rank_store_bounded_when_survivors_spread`: every batch keeps two rows forever; before, the store pinned all of every batch, and gathering without compaction would still pin `~B²` rows. - `test_partitioned_topk_dense_rank_runs_under_bounded_memory_pool`: the same shape at 256 partitions × 256 batches against a `GreedyMemoryPool` of 80× one input batch, mirroring the ROW_NUMBER and RANK pool tests. The PR peaks at ~39× one batch and stays flat; on `main` the test fails at batch 40 of 256 with `ResourcesExhausted` (~301× needed). - `test_partitioned_topk_dense_rank_size_charges_partition_keys` subtracts both copies of the keys that `size()` also sees (the store's `pk` column and the partition-key scratch rows), so it fails if the interned keys go uncharged. - The existing store tests are reworked for gathered entries. - Each of these mutants fails at least one test: no compaction, an in-flight release sent to the store, `groups_bytes` never decremented, `live_slots` never decremented, the pre-filter dropping boundary-equal rows, a no-op `repoint`, table growth left uncharged, `groups_bytes` dropped from `size()`, and `index_bytes` dropped from `allocated_size()`. - `datafusion-physical-plan` lib tests and the `window*`, `window_topn` and `range_partitioning` sqllogictests pass. ## Are there any user-facing changes? No API change; the operator is `pub(crate)`. Queries using `DENSE_RANK` with `enable_window_topn = true` use memory proportional to the rows they keep rather than to their input. -- 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]
