[SPARK-58404][SQL] Support bypassing partial WindowGroupLimit

[ulysses-you]

What changes were proposed in this pull request?

This PR adds a new config spark.sql.execution.bypassPartialWindowGroupLimit (default false). When enabled, the WindowGroupLimit planner strategy emits only the final WindowGroupLimitExec after the shuffle and skips constructing the pre-shuffle partial WindowGroupLimitExec, provided the window has a non-empty partition spec.

The final node still declares its required child distribution (ClusteredDistribution(partitionSpec)), so EnsureRequirements inserts the shuffle regardless; dropping the partial node removes both the pre-shuffle local top-k filter and the local SortExec that the partial node's requiredChildOrdering forces via EnsureRequirements. Skipping that local sort is a good part of the win. This mirrors the shape of the existing partial/final split for window group limit and is analogous to bypassing partial aggregation.

The bypass is gated on a non-empty partition spec. With an empty partition spec the final node requires AllTuples, so the shuffle funnels the whole input into a single reducer, while the partial pass we would drop bounds each input partition to limit rank groups before that shuffle. Bypassing it there would shuffle all raw rows to a single partition and is strictly worse than the normal Partial+Final path, so the partial is always kept. This mirrors the grouping-keys carve-out in bypassPartialAggregation (AggUtils.planAggregateWithoutDistinct).

Why are the changes needed?

The pre-shuffle partial WindowGroupLimitExec only pays off when it actually reduces rows. When each pre-shuffle partition already has few rows per window-partition group (a low reduction ratio), the partial node adds sorting/iteration cost with little benefit. This config lets users bypass the partial phase for such workloads.

Does this PR introduce any user-facing change?

Yes. A new config spark.sql.execution.bypassPartialWindowGroupLimit is added, defaulting to false, which preserves the existing behavior. It is also documented in the SQL performance tuning guide.

How was this patch tested?

Added a unit test in DataFrameWindowFunctionsSuite covering row_number/rank/dense_rank. For a partitioned window it asserts the executed plan contains a single WindowGroupLimitExec when the config is on and two when it is off; for an unpartitioned window it asserts both nodes remain regardless of the config, since the bypass is gated off there. Results are unchanged across all settings.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)

[uros-b]

Thank you @ulysses-you! I left a few comments, PTAL

[uros-b]

The test exercises only row_number(), but the PR description says the bypass applies to all three rank-like functions (row_number, rank, dense_rank), yet rank() and dense_rank() are not tested with bypassPartialWindowGroupLimit=true. Please add checks for rank() and dense_rank() too, this will also guard against future regressions.

[uros-b]

Also, the test covers only a non-empty partitionSpec. When partitionSpec is empty the Final node requires AllTuples distribution; bypassing partial is arguably more beneficial here (the partial pass cannot reduce cardinality across partitions, so it is pure overhead). A test asserting that bypassPartialWindowGroupLimit=true with an unpartitioned window produces a single WindowGroupLimitExec and correct results would make this boundary explicit.

[peter-toth]

Finding 1 (Blocking) and finding 2 (Non-blocking), both on this point -- I read the empty-partitionSpec case the other way round, so flagging it here rather than opening a new thread.

The partial pass never reduces cardinality across partitions -- not here, and not with a partition spec either. It's a mapPartitions filter, so all it ever does is prune within one input partition. And that pruning is sound because a row's global rank is always >= its rank inside a single input partition: anything with global rank <= limit also has local rank <= limit, so the partial can't drop a row the final needs. What it gives you is a hard output bound of limit rank groups per input partition, independent of the data.

That bound is why the empty-partitionSpec case is where the partial is worth most, not least:

-- bypass off
WindowGroupLimitExec Final          <- AllTuples
+- Sort [o]
   +- Exchange SinglePartition      <- carries ~ numPartitions * limit rank groups
      +- WindowGroupLimitExec Partial
         +- Sort [o]                <- one local sort per input partition, in parallel
            +- Scan t

-- bypass on
WindowGroupLimitExec Final
+- Sort [o]                         <- sorts the whole table, one task
   +- Exchange SinglePartition      <- carries the whole table
      +- Scan t

So the flag here costs a full-table shuffle into one reducer plus a single-threaded sort of everything, and buys back only the parallel local sorts it skipped. Downside unbounded, upside bounded.

With a partition spec it's the opposite: the exchange is a hash exchange to many reducers so post-shuffle work stays parallel, and the partial's bound is limit rank groups per input partition per key -- many keys with few rows each means it prunes nothing and is pure overhead. That is the "low reduction ratio" workload the config is for, and it can only really happen when there is a partition spec.

This is the same asymmetry bypassPartialAggregation ran into -- AggUtils.scala:136-141:

The bypass is only beneficial when there are grouping keys (groupingExpressions.nonEmpty): global aggregations (no GROUP BY) always produce a single output row, so the pre-shuffle partial aggregation achieves the maximum possible reduction ratio and should never be skipped. Bypassing a global aggregation would shuffle all raw rows to a single partition with no benefit, which is strictly worse than the normal Partial+Final path.

An empty partitionSpec is that same situation, so I'd mirror the guard in SparkStrategies.scala:808:

val finalChild = if (conf.bypassPartialWindowGroupLimit && partitionSpec.nonEmpty) {

If leaving it ungated is a deliberate "the user asked for it, don't second-guess" call, that's defensible, but then it's worth stating in the config doc.

Either way, the comment added at DataFrameWindowFunctionsSuite.scala:1723-1727 needs rewording -- it currently writes "the partial pass (which cannot reduce cardinality across partitions here)" into the test as the justification. Separately in that same comment, the row_number() note drops a condition: InferWindowGroupLimit.scala:118-122 only rewrites to Limit when partitionSpec is empty and limit < topKSortFallbackThreshold. It holds for limit = 1 here, so the test is fine, but the exclusion isn't unconditional.

[ulysses-you]

thank you @peter-toth , I'm fine to only apply on non-empty partitionSpec to align with agg side.

[ulysses-you]

thank you @uros-b for review, addressed comments

[ulysses-you]

also cc @cloud-fan @viirya @peter-toth thank you

[peter-toth]

Thanks for the PR, @ulysses-you!

The change is small and the shape reads correctly to me: the final WindowGroupLimitExec keeps declaring ClusteredDistribution(partitionSpec) / AllTuples, so EnsureRequirements still inserts the exchange, and the partial node is a pure pre-shuffle filter, so dropping it can't change results. My one substantive concern is the empty-partitionSpec case, which the new test explicitly turns on: there the partial is bounded to limit rank groups per input partition and the exchange funnels everything into a single reducer, so it's the one shape where the flag's downside has no ceiling. bypassPartialAggregation carves out the analogous global-aggregation case for exactly this reason (AggUtils.scala:136-141). The rest is a docs placement and a config-metadata detail.

@uros-b's two coverage asks are addressed by 6e5c511. Findings 1 and 2 below disagree with the reasoning in the second one, so I've put them as a reply on that thread rather than opening new ones.

Blocking

• 1. Bypass isn't gated on a non-empty partition spec: with an empty partitionSpec the final node requires AllTuples, so the bypass ships the whole input through one reducer and sorts it there, while the partial it drops is bounded to limit rank groups per input partition. Either mirror bypassPartialAggregation's carve-out or say in the config doc why window group limit differs. [reply on @uros-b's thread]

Non-blocking

• 2. Test comment states the rationale backwards: "the partial pass (which cannot reduce cardinality across partitions here)" -- the partial never reduces across partitions in either case; what it does is cut each input partition down, which matters more when the destination is a single reducer. [reply on @uros-b's thread]
• 3. Doc row landed in the wrong table: it's the last entry of the "Tuning Partitions" table, next to the parallel-partition-discovery configs. [inline: docs/sql-performance-tuning.md:140]
• 4. Binding policy should be NOT_APPLICABLE: the config is read in physical planning and can't change how a view body resolves, which is step 1 of the ConfigBindingPolicy decision procedure. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:4537]

Minor

• 5. Description understates what is dropped: removing the partial node also removes the pre-shuffle local SortExec that its requiredChildOrdering forces (EnsureRequirements.scala:285-297), not just "the pre-shuffle local top-k filter" -- worth saying, since that sort is a good part of the win.

[peter-toth]

Finding 3. This row is the last entry of the "Tuning Partitions" table (spark.sql.files.*, spark.sql.shuffle.partitions, spark.sql.sources.parallelPartitionDiscovery.*), which is about how input and shuffle partitions get sized. A window top-k planning flag isn't that, and someone looking for window tuning won't find it there.

The page is organised as topical sections that each own a table (## Tuning Partitions, ## Optimizing the Join Strategy, ## Storage Partition Join), so the cleanest fix is a short section of its own -- e.g. a ## Tuning Window Functions after ## Tuning Partitions, with a one-line intro and this row. That would also be the natural home for spark.sql.execution.bypassPartialAggregation, which is undocumented on this page today, though that's not this PR's job.

[ulysses-you]

addressed

[peter-toth]

Finding 4. Step 1 of the ConfigBindingPolicy decision procedure is "Can the config change the result of resolving the body of a view/UDF/procedure, i.e. the resolved plan? If not, use NOT_APPLICABLE", with the explicit note that "even physical planning or runtime configs may be read while resolving a view, but they do not change what the body resolves to". This one is read in SparkStrategies, after the body is resolved and optimized, so NOT_APPLICABLE looks like the right declaration -- as it is for USE_HASH_AGG (:4091, also a physical-planning operator choice) and for WINDOW_SEGMENT_TREE_ENABLED just below (:4543).

Behaviour is the same either way (NOT_APPLICABLE also reads from the active session), so this is purely about the declaration being right. I assume it was copied from BYPASS_PARTIAL_AGGREGATION (:4103), which by the same reading looks mis-declared too -- separate cleanup, not this PR.

Suggested change

	.withBindingPolicy(ConfigBindingPolicy.SESSION)
	.withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)

[ulysses-you]

addressed

[peter-toth]

Re-checked through c3ded05 -- findings 1-5 all resolved: the bypass is gated on partitionSpec.nonEmpty (SparkStrategies.scala:816), the test comment now states the rationale the right way round and names the topKSortFallbackThreshold condition on the row_number() exclusion, the doc row moved into its own ## Tuning Window Functions section, the binding policy is NOT_APPLICABLE, and the description names the pre-shuffle local SortExec that goes away with the partial node. Nothing regressed. One new item.

Non-blocking

• 6. No measurement behind a performance-only flag (late catch): the PR's whole case is a perf win and there are no numbers, while TopKBenchmark already benchmarks this exact query shape with checked-in results. A bypass case there (or numbers in the description) would show at which reduction ratios the flag pays off. [inline: docs/sql-performance-tuning.md:186]

[peter-toth]

Finding 6. This is the claim the whole PR rests on, and nothing in the PR measures it. sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/TopKBenchmark.scala already benchmarks this exact shape -- 20M rows in 11 partitions, PARTITION BY b and unpartitioned, all three rank-like functions, WINDOW_GROUP_LIMIT_THRESHOLD on and off -- with checked-in results under sql/core/benchmarks/TopKBenchmark-results.txt, so the harness is free.

Worth noting that its partitioned case sits at the wrong end of the range for this flag: id % 1024 as b over 20M rows in 11 partitions is ~1.8k rows per key per input partition against limit = 200, so the partial prunes ~90% and the bypass should lose there. The workload the config targets is the opposite -- more keys than the partial can prune, i.e. fewer rows per key per input partition than limit:

spark.range(0, N, 1, 11).selectExpr("id as a", "id % 1024 as b", "id % 4000000 as c")
...
Seq("PARTITION BY b", "PARTITION BY c").foreach { partition =>
  Seq(false, true).foreach { bypass =>
    benchmark.addCase(s"$function ($partition, bypassPartial: $bypass)") { _ =>
      withSQLConf(BYPASS_PARTIAL_WINDOW_GROUP_LIMIT.key -> bypass.toString) {
        f(function, partition)
      }
    }
  }
}

A pair of numbers from each end answers the question a reader of this config doc actually has -- "is my reduction ratio low enough" -- and shows whether skipping the pre-shuffle sort pays for the bigger shuffle when the partial does prune.

[cloud-fan]

1 blocking, 2 non-blocking, 0 nits.
The physical-plan rewrite is sound, but the session-scoped public API does not match this query-specific tuning decision and should be replaced with a query hint.

Already raised in existing discussion (1)

• The new user-facing flag is justified solely as a performance optimization, but this patch provides no measurements showing when bypassing the partial phase wins or loses. Add a low-reduction-ratio case to TopKBenchmark (or equivalent numbers to the PR description) so users and maintainers can validate the claimed tradeoff. -- existing discussion

Design / architecture (1)

• sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:4531: Replace the public session configuration with a query-level hint. Whether Partial pays off depends on each query's limit and per-key distribution, so one session-wide value can help one window query while regressing another under the same SQLConf. Spark already represents hints as logical-plan nodes scoped to their child; carry a dedicated hint to the generated logical WindowGroupLimit and test its scope when a query contains multiple windows. -- see inline

Suggestions (1)

• sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala:812: The empty-partition rationale is stated as universal, but an already-single-partition child satisfies AllTuples without a shuffle. In that case the partial pass cannot reduce shuffle input and only adds another WindowGroupLimit pass. Keep the guard, but narrow the comment and config documentation to the multi-partition case so the code does not claim that retaining Partial always helps or that bypassing it is always strictly worse. -- see inline

Verification

I traced the strategy through WindowGroupLimitExec and EnsureRequirements: Final preserves output and continues to enforce distribution and ordering, so removing Partial does not alter results. I also checked Spark's hint representation, which attaches an UnresolvedHint to a specific logical-plan child, providing the query-local control surface this optimization needs.

PR description suggestions

• Explain the user-control scope: pruning efficiency varies per query, so justify why a session configuration is appropriate or describe the query-hint contract that replaces it.

[cloud-fan]

Please make this a query-level hint instead of a public session config. The payoff depends on each query's limit and per-key distribution, so one session value can improve one window query while regressing another; hints already provide a logical-plan-local control surface (UnresolvedHint carries its child). Carry a dedicated hint to the generated WindowGroupLimit, define its behavior when a query has multiple windows, and cover that scope in tests.

[ulysses-you]

thank you @cloud-fan for the suggestion, yeah, it is good if we can have query-level hint for finer-grained control, but I think a seesion config is still useful:

1. it already covers the common case well
2. it is easy to use for user, hint need to touch the query or even code

[cloud-fan]

A conf to turn on/off PartialWindowGroupLimit for the entire session does not make sense, or in your case every query has a dedicated session?

[ulysses-you]

@cloud-fan this is the truth.. most of our SQL queries have its own driver and session, so it is a kind of query level conf.

[cloud-fan]

Spark is a general engine, we should not add a change that only makes sense for certain env... In you case, you need to touch every one-query-application to add this conf, it seems the same effort to add the query hint.

[ulysses-you]

yes, I agree we need to touch every one-query-application, but adjusting configurations is not exactly the same as tuning SQL queries. Tuning SQL queries is usually more troublesome, and we still have many JAR tasks that need to be recompiled and repackaged..

[cloud-fan]

You can add this conf in your spark fork. This is not a feature rollout conf, as we clearly can not bypass partial window group limit always.

We should either use query hint, or build a real optimization that only skip partial window group limit when it's beneficial (e.g. stats-based decision, or runtime adaptivity).

[cloud-fan]

Please narrow this claim to multi-partition inputs. An already-single-partition child satisfies AllTuples without a shuffle (EnsureRequirements.scala:71), so Partial cannot reduce shuffle input there and only adds another WindowGroupLimit pass. The guard is still a sensible default, but always helps and strictly worse overstate its rationale; the config docs make the same unconditional assumption.

[ulysses-you]

addressed

[dongjoon-hyun]

+1, LGTM. Thank you, @ulysses-you and all.

[dongjoon-hyun]

Merge Summary:

• merged into master 91c03e3
• merged into branch-4.x 8af2f3d
• merged into branch-4.3 185f0da

Posted by merge_spark_pr.py

[HyukjinKwon]

1 blocking, 1 non-blocking, 0 nits.
The physical-plan rewrite is correct and well-guarded; the open question is whether a session config is the right mechanism for a per-query tuning decision.

Already raised in existing discussion (1)

• A performance-only flag with no measurements. TopKBenchmark already benchmarks this exact query shape with checked-in results, so adding a bypass case there (or numbers in the description) would show at which reduction ratios the bypass actually pays off -- and would directly inform the config-vs-hint discussion above. -- existing discussion

Design / architecture (1)

• General: The bypass is exposed as a session-scoped config, but whether skipping the partial WindowGroupLimit helps is a per-query property (it only wins at a low pre-shuffle reduction ratio). A committer (cloud-fan) has objected on the PR that this should be a query hint or a stats/AQE-based runtime decision rather than a session flag, and that discussion is unresolved. The code itself is correct (results are unchanged; the final node still forces the shuffle), so this is a design/API direction question, not a bug. Worth surfacing that the session-config vs query-hint disagreement is still open and likely blocks merge in its current form.

Verification

Confirmed results are unchanged (final node forces the shuffle; partial is a pure pre-shuffle filter) and the empty-partitionSpec carve-out matches the bypassPartialAggregation precedent. The earlier reviewer findings (1-5: gating, test comment, doc placement, binding policy, description) are all resolved in the current source. Two items remain open on the PR and are reflected below.

[cloud-fan]

@dongjoon-hyun I raised a blocking issue and the author has not addressed it, see #57602 (comment) , why we merge it?

[dongjoon-hyun]

Oh, I thought all review comments addressed. Do you want me revert this PR, @cloud-fan ? If you want, I'll do right now.

[cloud-fan]

yea let's revert first. Further more, @ulysses-you has a new PR that can adaptively bypass partial aggregate: #57742 . I think the same idea can apply to partial window group limit as well, so this PR may be closed eventually.

[dongjoon-hyun]

Sure, sorry for the trouble, @cloud-fan .

[dongjoon-hyun]

This is reverted.

• branch-4.3: 7c60866
• branch-4.x: 88b4c94
• master: c65076b