Skip to content

feat(window-topn): streaming per-partition top-K for pre-sorted input - #24629

Closed
SubhamSinghal wants to merge 3 commits into
apache:mainfrom
SubhamSinghal:window-topn-streaming
Closed

SubhamSinghal wants to merge 3 commits into
apache:mainfrom
SubhamSinghal:window-topn-streaming

Conversation

@SubhamSinghal

@SubhamSinghal SubhamSinghal commented Aug 24, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Relates to the WindowTopN optimization (ROW_NUMBER/RANK/DENSE_RANK per-partition top-K). It adds the missing already-sorted-input path.
related to #24060

Rationale for this change

The WindowTopN rule today only handles the case where the window's input is unsorted: it replaces the SortExec with a heap-based PartitionedTopKExec (one top-K heap per distinct partition key). That is the right trade-off when a sort would otherwise be needed, but it has two costs that make it ship default-off:

  • O(K × P) memory (P = number of distinct partitions), which regresses at high partition cardinality.
  • Blocking emission — it buffers all input before emitting.

When the window's input is already sorted by (partition_keys, order_keys) — i.e. EnforceSorting elided the SortExec because the ordering was already satisfied — the heaps are pure waste: the data is in the exact order we need, so we can compute the per-partition rank in a single streaming pass with O(1) state and emit incrementally, dropping rows whose rank exceeds K before they reach the (expensive) full window.

Today that already-sorted case doesn't fire at all: the rule requires a SortExec child, so the plan falls through to the default full-window-then-filter. This PR closes that gap.

What changes are included in this PR?

A new sibling operator and a restructured rule with two cases:

  • StreamingPartitionedTopKExec (sorts/streaming_partitioned_topk.rs) — requires input sorted by (partition_keys, order_keys), computes the per-partition rank in one pass (carrying only prev_partition / prev_order / rank / count), and emits filtered batches incrementally. EmissionType::Incremental, maintains_input_order = true, required_input_ordering = (partition, order), KeyPartitioned distribution. Supports ROW_NUMBER, RANK and DENSE_RANK (matching WindowFnKind);.
  • WindowTopN rule split into a shared matcher plus two classifiers:
    • try_streaming — input already sorted (no SortExec): insert the streaming operator between the window and its child. Applied unconditionally, because it is a pure win (cull-before-window, O(1) memory, order-preserving).
    • try_heap — input unsorted (SortExec present): today's behavior, replace the SortExec with PartitionedTopKExec. Still gated on enable_window_topn.

Are these changes tested?

Yes

Are there any user-facing changes?

New physical operator StreamingPartitionedTopKExec in EXPLAIN output for qualifying queries over pre-sorted input. Results are unchanged. Because the streaming rewrite is a pure win, it applies even when enable_window_topn is disabled (that flag continues to gate only the heap path).

@github-actions github-actions Bot added documentation Improvements or additions to documentation optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) common Related to common crate physical-plan Changes to the physical-plan crate labels Aug 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate core Core DataFusion crate documentation Improvements or additions to documentation optimizer Optimizer rules physical-plan Changes to the physical-plan crate sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant