Repository navigation
perf: share encoder/reservation across PartitionedTopKExec partition … - #23096
Conversation
|
@kosiew can you review this PR? |
There was a problem hiding this comment.
@SubhamSinghal
Thanks for the change! The implementation and the approach looks good to me. I just have a couple of small, non-blocking suggestions.
| /// Builds a `(pk Int32, val Int32)` schema and a `PartitionedTopK` | ||
| /// partitioned by `pk` with order `val ASC`. Helper for the | ||
| /// `PartitionedTopK` tests below. | ||
| fn build_partitioned_topk(k: usize) -> Result<(Arc<Schema>, PartitionedTopK)> { |
There was a problem hiding this comment.
Nice addition to the test coverage. One thing that might be worth adding is a regression covering ORDER BY val DESC and/or NULL sort values. This path relies on Arrow row encoding for sort direction and null ordering, so having one extra SQL-level or unit test would help protect the shared encoder rewrite across the main ordering variants used by window Top-N queries. Not blocking since this reuses the existing build_sort_fields/RowConverter path and the current targeted tests pass.
There was a problem hiding this comment.
Added UT covering null first, null last and sort desc direction
| SELECT id3 % 100000 AS pk, v2 AS largest2_v2, | ||
| ROW_NUMBER() OVER (PARTITION BY id3 % 100000 ORDER BY v2 DESC) AS order_v2 | ||
| FROM large WHERE v2 IS NOT NULL | ||
| ) sub_query WHERE order_v2 <= 2; No newline at end of file |
There was a problem hiding this comment.
Small formatting nit: could you add a trailing newline at EOF? No behavioral impact, just keeps the file consistent.
## Which issue does this PR close? - Related to apache#6899 (DENSE_RANK to follow in a separate PR will close it). ## Rationale for this change PR apache#21479 introduced `WindowTopN` for `ROW_NUMBER` only; `RANK` and `DENSE_RANK` were explicitly out of scope. This PR extends the rule to `RANK`, replacing the full sort under `Filter(rk≤K) → Window(RANK) → Sort` with a per-partition heap-of-K plus a boundary-tie buffer. ## What changes are included in this PR? - **`datafusion/physical-plan/src/topk/mod.rs`** — new `pub(crate) struct PartitionedTopKRank` (sibling of `PartitionedTopK` from apache#23096) with per-partition `RankPartitionState { TopKHeap, Vec<TieEntry> }`. - **`datafusion/physical-plan/src/sorts/partitioned_topk.rs`** — `WindowFnKind` enum (`RowNumber` / `Rank`). `do_partitioned_topk` dispatches on `fn_kind` to `PartitionedTopK::try_new` or `PartitionedTopKRank::try_new`; - **`datafusion/physical-optimizer/src/window_topn.rs`** — `is_row_number` → `supported_window_fn(expr) -> Option<WindowFnKind>`; empty-`order_by` guard for RANK; `WindowFnKind` plumbed through `PartitionedTopKExec::try_new`. - **`datafusion/sqllogictest/test_files/window_topn.slt`** — RANK SLT cases: basic, strict (`<`), flipped (`>=` / `>`), boundary ties, ties spanning ob values, empty-`ORDER BY` (rule must NOT fire), mixed window functions, ASC/DESC × NULLS FIRST/LAST, QUALIFY. - **`datafusion/core/tests/physical_optimizer/window_topn.rs`** — 6 new RANK rule unit tests covering predicate matching, partition-by/order-by guards, dense_rank skip. - **`benchmarks/queries/h2o/window.sql`** — six new RANK queries (Q14–Q17, Q22, Q23) covering partition counts from ~100 to ~100K, low and heavy tie densities. h2o `window` benchmark, 10M-row `large` table, RANK top-2, 3-iteration average. Toggle via `DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN`. | Variant | Partitions | OFF (rule disabled) | ON (rule enabled) | Δ | |---|---:|---:|---:|:---:| | RANK low ties (`id3 % 100`) | ~100 | 305 ms | **107 ms** | **2.84× faster** ✓ | | RANK low ties (`id3 % 1000`) | ~1K | 263 ms | **120 ms** | **2.19× faster** ✓ | | RANK heavy ties (`id3 % 1000`, `v2 % 10` OB) | ~1K | 282 ms | **125 ms** | **2.25× faster** ✓ | | RANK low ties (`id2`) | ~10K | 363 ms | **140 ms** | **2.59× faster** ✓ | | RANK heavy ties (`id2`, `v2 % 10` OB) | ~10K | 291 ms | **143 ms** | **2.04× faster** ✓ | | RANK low ties (`id3 % 100K`) | ~100K | 241 ms | 422 ms | 1.75× slower | ## Are these changes tested? Yes: - `cargo test -p datafusion-physical-plan --lib` — 1455 passed - `cargo test -p datafusion-physical-optimizer --lib` — 27 passed - `cargo test -p datafusion --test core_integration physical_optimizer::window_topn::` — 13 passed (7 ROW_NUMBER + 6 RANK) - `cargo test --test sqllogictests -- window_topn` — passed ## Are there any user-facing changes? The existing `optimizer.enable_window_topn` config flag (default `false`) now also covers `RANK` queries. No public API additions
## Which issue does this PR close? - Related to apache#6899 (DENSE_RANK to follow in a separate PR will close it). ## Rationale for this change PR apache#21479 introduced `WindowTopN` for `ROW_NUMBER` only; `RANK` and `DENSE_RANK` were explicitly out of scope. This PR extends the rule to `RANK`, replacing the full sort under `Filter(rk≤K) → Window(RANK) → Sort` with a per-partition heap-of-K plus a boundary-tie buffer. ## What changes are included in this PR? - **`datafusion/physical-plan/src/topk/mod.rs`** — new `pub(crate) struct PartitionedTopKRank` (sibling of `PartitionedTopK` from apache#23096) with per-partition `RankPartitionState { TopKHeap, Vec<TieEntry> }`. - **`datafusion/physical-plan/src/sorts/partitioned_topk.rs`** — `WindowFnKind` enum (`RowNumber` / `Rank`). `do_partitioned_topk` dispatches on `fn_kind` to `PartitionedTopK::try_new` or `PartitionedTopKRank::try_new`; - **`datafusion/physical-optimizer/src/window_topn.rs`** — `is_row_number` → `supported_window_fn(expr) -> Option<WindowFnKind>`; empty-`order_by` guard for RANK; `WindowFnKind` plumbed through `PartitionedTopKExec::try_new`. - **`datafusion/sqllogictest/test_files/window_topn.slt`** — RANK SLT cases: basic, strict (`<`), flipped (`>=` / `>`), boundary ties, ties spanning ob values, empty-`ORDER BY` (rule must NOT fire), mixed window functions, ASC/DESC × NULLS FIRST/LAST, QUALIFY. - **`datafusion/core/tests/physical_optimizer/window_topn.rs`** — 6 new RANK rule unit tests covering predicate matching, partition-by/order-by guards, dense_rank skip. - **`benchmarks/queries/h2o/window.sql`** — six new RANK queries (Q14–Q17, Q22, Q23) covering partition counts from ~100 to ~100K, low and heavy tie densities. h2o `window` benchmark, 10M-row `large` table, RANK top-2, 3-iteration average. Toggle via `DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN`. | Variant | Partitions | OFF (rule disabled) | ON (rule enabled) | Δ | |---|---:|---:|---:|:---:| | RANK low ties (`id3 % 100`) | ~100 | 305 ms | **107 ms** | **2.84× faster** ✓ | | RANK low ties (`id3 % 1000`) | ~1K | 263 ms | **120 ms** | **2.19× faster** ✓ | | RANK heavy ties (`id3 % 1000`, `v2 % 10` OB) | ~1K | 282 ms | **125 ms** | **2.25× faster** ✓ | | RANK low ties (`id2`) | ~10K | 363 ms | **140 ms** | **2.59× faster** ✓ | | RANK heavy ties (`id2`, `v2 % 10` OB) | ~10K | 291 ms | **143 ms** | **2.04× faster** ✓ | | RANK low ties (`id3 % 100K`) | ~100K | 241 ms | 422 ms | 1.75× slower | ## Are these changes tested? Yes: - `cargo test -p datafusion-physical-plan --lib` — 1455 passed - `cargo test -p datafusion-physical-optimizer --lib` — 27 passed - `cargo test -p datafusion --test core_integration physical_optimizer::window_topn::` — 13 passed (7 ROW_NUMBER + 6 RANK) - `cargo test --test sqllogictests -- window_topn` — passed ## Are there any user-facing changes? The existing `optimizer.enable_window_topn` config flag (default `false`) now also covers `RANK` queries. No public API additions
…he#27197) ## Which issue does this PR close? - Closes apache#24608. ## Rationale for this change `enable_window_topn` has been off since apache#21479 because of high-cardinality regressions against the sort plan. The work items in apache#24608 are now done: - apache#23096, - apache#25730, - apache#25968, - apache#26133 The h2o window benchmark's Top-N queries (`benchmarks/queries/h2o/window.sql`, Q13–Q29, small dataset, `J1_1e7_1e7_NA.parquet`) are now faster with the flag on in every case, up to 100K partitions and including the heavy-ties variants. `dfbench h2o --iterations 5`, median, `DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN=false` for the baseline, same binary, Apple M4 Max: | Query | function, partitions | off (ms) | on (ms) | change | |---|---|---:|---:|---| | Q13 | ROW_NUMBER, id2 | 127.4 | 25.8 | 4.93x faster | | Q14 | ROW_NUMBER, 100 | 190.8 | 35.9 | 5.32x faster | | Q15 | ROW_NUMBER, 1K | 173.6 | 35.7 | 4.86x faster | | Q16 | ROW_NUMBER, 10K | 149.7 | 33.0 | 4.54x faster | | Q17 | ROW_NUMBER, 100K | 151.1 | 49.1 | 3.08x faster | | Q18 | RANK, 100 | 226.3 | 28.8 | 7.86x faster | | Q19 | RANK, 1K | 186.0 | 29.6 | 6.28x faster | | Q20 | RANK, 1K, heavy ties | 209.8 | 57.1 | 3.67x faster | | Q21 | RANK, id2 | 174.2 | 25.4 | 6.87x faster | | Q22 | RANK, id2, heavy ties | 190.5 | 37.2 | 5.13x faster | | Q23 | RANK, 100K | 188.1 | 63.5 | 2.96x faster | | Q24 | DENSE_RANK, 100 | 248.4 | 38.6 | 6.43x faster | | Q25 | DENSE_RANK, 1K | 199.6 | 36.9 | 5.40x faster | | Q26 | DENSE_RANK, 1K, heavy ties | 232.2 | 43.3 | 5.37x faster | | Q27 | DENSE_RANK, id2 | 211.3 | 45.7 | 4.63x faster | | Q28 | DENSE_RANK, id2, heavy ties | 208.7 | 53.6 | 3.89x faster | | Q29 | DENSE_RANK, 100K | 205.8 | 123.3 | 1.67x faster | ## What changes are included in this PR? - `enable_window_topn` defaults to `true`; its doc comment now lists ROW_NUMBER, RANK and DENSE_RANK and drops the high-cardinality warning. - `window_topn.slt` restores the setting with `RESET` instead of `SET ... = false`. - `configs.md` and `information_schema.slt` regenerated. - The h2o `window_sorted` settings comment no longer says the flag defaults to false (it still sets it explicitly). ## What is the testing strategy for this PR? Existing coverage. With the flag on by default, the whole sqllogictest suite and `core_integration` pass with only the `information_schema.slt` default change; `qualify.slt`'s `rank <= 3` / `dense_rank <= 2` cases now run through `PartitionedTopKExec` and return the same results. `window_topn.slt` still covers the disabled path explicitly. ## Are there any user-facing changes? Yes. Top-K-per-partition queries (`rn <= K`, `rn < K` over ROW_NUMBER / RANK / DENSE_RANK) plan with `PartitionedTopKExec` instead of a full sort, so `EXPLAIN` output changes; results don't. Set `datafusion.optimizer.enable_window_topn = false` to get the old plan. --------- Co-authored-by: Andrew Lamb <[email protected]>
Which issue does this PR close?
Follow-up to #21479 (
PartitionedTopKExecforROW_NUMBER ... PARTITION BY ... LIMIT N) toward closing #6899.Rationale for this change
PartitionedTopKExectoday maintains aHashMap<OwnedRow, TopK>— one fullTopKper distinct partition key. EachTopKcarries its ownRowConverter,MemoryReservationregistered with the runtime pool,TopKMetrics, and scratchRowsbuffer. With high partition cardinality every partition seen for the first time pays:RowConverter::new(parsesSortFieldlist, allocates per-encoder state)MemoryConsumer::registerwith the pool (involves a global lock)TopKMetricssetupRows::empty_rowsallocationFor the h2o window-TopN sweep on a 10M-row CSV (
id3 % Npartition cardinality), this shows up as a regression at ≥10K partitions —PartitionedTopKExecis slower than the unpartitionedSortExecbaseline that it's meant to replace.What changes are included in this PR?
Adds a
PartitionedTopKsibling type totopk/mod.rsthat holds the shared encoder/reservation/metrics state once at the operator level and aHashMap<OwnedRow, TopKHeap>of cheap per-partition heap state.PartitionedTopKExecswitches fromHashMap<OwnedRow, TopK>to onePartitionedTopK.Bench results
Today's default (main, flag-off) vs this PR (flag-on)
h2o
id3 % Nsweep, 10M-row CSV, 3 iterations per query, release build,enable_window_topn=trueon both sides:main10K is the inflection point: on
mainit's a regression vs the sort baseline (640 ms vs 234 ms); after this PR it's a win (137 ms vs 234 ms — 1.7× faster than sort). 100K nearly catches up to the sort baseline (320 ms vs 238 ms).enable_window_topndefault staysfalseper the #21479 discussion — 100K+ remains slower than sort on average, so this PR doesn't motivate flipping the default. It's the prerequisite for further optimizations that would attack the residual 100K+ cliff.Are these changes tested?
Yes
Are there any user-facing changes?
No public API changes.