[SPARK-58233][SQL] Push down local sort to reduce redundant sorts within a stage

[ulysses-you]

What changes were proposed in this pull request?

Adds a new physical rule PushDownLocalSort, wired into both the AQE queryStagePreparationRules and the non-AQE QueryExecution preparations (right after EnsureRequirements, before CombineAdjacentAggregation), gated by a new internal config spark.sql.execution.pushDownLocalSort (default true).

EnsureRequirements inserts one local SortExec(global = false) above every operator whose requiredChildOrdering is not satisfied. When two such requirements are in a prefix-cover relationship, a stage computes multiple local sorts that only differ in width. The rule pushes the wider local sort down through order-preserving operators onto the narrower one below, widening it so a single sort serves both operators and the redundant upper sort is dropped.

For example, a sort aggregate stacked on a window over the same clustering keys:

SortAggregate(key = [a, b, c])
  Sort([a, b, c], global = false)   <- upper, wider
    Window([a], [b])
      Sort([a, b], global = false)  <- lower, narrower
        Exchange(hashpartitioning([a]))

becomes:

SortAggregate(key = [a, b, c])
  Window([a], [b])                  requiredChildOrdering [a, b] satisfied by [a, b, c]
    Sort([a, b, c], global = false) <- single sort now serves both operators
      Exchange(hashpartitioning([a]))

Order-preserving operators traversed: ProjectExec, WindowExecBase, and — only when spark.sql.execution.pushDownLocalSort.throughCardinalityReducer is enabled (default false) — the cardinality reducers FilterExec and WindowGroupLimitExec. Crossing a reducer can move the wider sort from the surviving rows to the full input, which may outweigh the saved sort, so it is opt-in. A ProjectExec or FilterExec is crossed only when its expressions are deterministic (mirroring EliminateSorts.canEliminateSort), so a seeded non-deterministic expression is never re-associated with different rows. When a ProjectExec renames an ordering column in its output (b AS x), the ordering is rewritten from the project's output space back to its child's space as it is pushed through (plain renames only). The rule never crosses a shuffle or a non-order-preserving operator.

CollectMetricsExec is intentionally excluded: it observes rows in input order, so widening a sort under it would silently change order-sensitive observed metrics (first/last/collect_list).

This is complementary to RemoveRedundantSorts, which only removes a sort already satisfied by its child; it never moves a wider sort down to absorb a narrower one.

Why are the changes needed?

Stacked windows, or a sort aggregate over a window with prefix-compatible ordering requirements, compute several local sorts that a single wider sort could satisfy. Reusing the widest sort saves the redundant sorting work in the stage.

Does this PR introduce any user-facing change?

No. It only removes redundant local sorts from physical plans; the result of a deterministic query is unchanged. As with any plan change that alters within-partition row order, an order-sensitive window function over a non-unique ORDER BY (e.g. first_value/last_value/collect_list) may return a different — but still valid — non-deterministic result. Guarded by spark.sql.execution.pushDownLocalSort.

How was this patch tested?

New PushDownLocalSortSuite (AQE on/off variants) with:

• positive SQL cases: stacked windows, filter above the sorts, window to sort aggregate, renaming project;
• negative SQL cases: disjoint orderings, a shuffle between the sorts, sort-direction mismatch, a project that computes the ordering column;
• plan-level cases: multi-operator chain, renaming-project rewrite, three stacked sorts in one pass, non-deterministic project/filter are not crossed, and a cardinality-reducing filter is crossed only when the reducer config is enabled.

Verified the TPCDS/TPCH PlanStability golden plans are unchanged.

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

Generated-by: Claude Code (Opus 4.8)

[ulysses-you]

cc @viirya @cloud-fan @peter-toth @sunchao if you have time to take a look, thank you

[uros-b]

Thank you @ulysses-you!

[peter-toth]

Thanks for the PR, @ulysses-you!

I walked the rule independently: it fires on a local SortExec, walks down through order-preserving unary operators, and when it finds a lower local sort whose ordering the upper one strictly prefix-covers, widens the lower sort and drops the upper one, rewriting the ordering back through plain renames on the way. The orderingSatisfies + subsetOf(child.outputSet) guards look right, the re-exposed ordering matches for both the non-renaming and plain-rename paths, and the plan-level tests cover the multi-op / rename / three-stack cases. Nothing here produces wrong results.

The one thing worth a closer look: widening the sort that feeds a Window/WindowGroupLimit refines the within-tie input order, which changes the (already non-deterministic) result of order-sensitive window functions. That's the same order-sensitivity you cite for excluding CollectMetricsExec, so the "query results are unchanged" claim and the exclusion set feel inconsistent. All non-blocking.

Non-blocking

1. Order-sensitive window functions see a refined input order. Widening the sort below a Window/WindowGroupLimitExec changes the tie-break order the window sees, so collect_list/first_value/last_value/nth_value/lead/lag and ROWS-frame aggregates can return a different (valid, non-deterministic) result, and a row_number-based window group limit can keep a different set of rows. Within Spark's non-determinism contract, but it contradicts "query results are unchanged" and is the same reason you excluded CollectMetricsExec.
2. The window-feeds-aggregate test can't observe a tie-order change. The data makes c constant within every (a, b) group, so widening [a,b]->[a,b,c] leaves the window's input order unchanged and the result comparison passes trivially -- it doesn't exercise finding 1.
3. SortAggregateExec handling looks unreachable and is untested. I couldn't build a plan where traversing a SortAggregateExec enables a widening; the guards seem to force upperOrder no wider than the grouping-key sort EnsureRequirements puts directly under it. If you have a firing case, add a test (also for the resultExpressions rename and WindowGroupLimitExec); otherwise consider dropping it.

[peter-toth]

Widening the sort that feeds a window changes the row order the window sees within ties of its own ORDER BY. For order-sensitive window functions that's observable:

• collect_list/collect_set/first_value/last_value/nth_value collect/pick in input order;
• lead/lag and aggregates with a ROWS frame depend on physical adjacency;
• a row_number()-based WindowGroupLimitExec keeps a different set of rows when ties reorder.

Example: collect_list(c) OVER (PARTITION BY a ORDER BY b) stacked under a wider window that needs [a, b, c]. Today the collect_list window sorts by [a, b]; after this rule it sorts by [a, b, c], so within equal b the list follows c.

These are all non-deterministic over a non-unique ORDER BY, so it isn't wrong per Spark's contract -- but it's exactly the order-sensitivity you use to keep CollectMetricsExec out of isOrderPreserving (and the reason cited in the PR description). So the description's "query results are unchanged" holds for deterministic queries but not for these, and the line between "exclude CollectMetrics" and "push through a window carrying first/last/collect_list" is hard to defend.

Two asks: (1) soften that claim in the description; (2) decide whether a window carrying an order-sensitive function should be treated like CollectMetricsExec. If you deliberately keep them in (they're non-deterministic anyway), a one-line comment here explaining why the window case is acceptable while CollectMetrics isn't would settle it for the next reader.

[ulysses-you]

Good point, thanks. I kept pushing through windows on purpose: it only refines the tie order for a non-unique window ORDER BY, which is already non-deterministic per Spark's contract, so no deterministic result changes. That is genuinely different from CollectMetricsExec, whose observed metric is an expected-stable side output rather than a non-deterministic query result — so I keep that one excluded.

I added a comment at the WindowExecBase/WindowGroupLimitExec entries in isOrderPreserving spelling out exactly this distinction, and softened the PR description: "query results are unchanged" -> "the result of a deterministic query is unchanged", plus an explicit note that order-sensitive window functions over a non-unique ORDER BY may return a different (still valid) non-deterministic result.

[peter-toth]

Good case, but the data neutralizes the risk it looks like it's testing. Within any (a, b) group every row has the same c: ids sharing (id%10, id%7) differ by multiples of lcm(10,7)=70, and 70 % 5 == 0, so id % 5 is constant per (a, b). Widening the window's sort from [a, b] to [a, b, c] therefore can't reorder within-(a,b) ties, so row_number comes out identical with the rule on and off and checkAnswer passes by construction.

If c actually varied within (a, b), the [a,b,c] sort would reorder ties and row_number would differ between the two configs -- which would make this comparison flaky rather than catch a bug. So this test can't observe the order change from the window comment. Worth either a note here, or a companion test that asserts the sort count drops while using a tie-safe function (rank/dense_rank) or a unique key, so the result stays stable even when ties genuinely reorder.

[ulysses-you]

You are right — I confirmed count(distinct c) per (a, b) is 1 for that data (lcm(10,7)=70, 70 % 5 == 0), so the widening could not reorder any tie and the check passed trivially. Fixed: switched to id % 13 for c, which varies within (a, b) (70 % 13 != 0), and changed the window function to the tie-safe RANK(). Now the [a,b]->[a,b,c] widening genuinely reorders rows within the window's ORDER BY b ties, while RANK() keeps the result stable — so the test exercises the reorder and still verifies correctness plus the sort-count drop. Added a comment explaining the data choice.

[peter-toth]

I couldn't construct a plan where traversing a SortAggregateExec here actually enables a widening. To push upperOrder through it, line 118 requires it to satisfy the grouping-key ordering and line 119 requires every key to be a child column -- together that forces upperOrder to be a prefix-cover of the grouping keys and nothing wider, because any extra output column is an aggregate result that fails the subsetOf(child.outputSet) check. But the aggregate already requires its input sorted by the full grouping keys, so EnsureRequirements places that sort directly beneath it; there's no strictly-narrower lower sort left for upperOrder to widen.

If that's right, the SortAggregateExec branch (here, plus the resultExpressions rewrite at line 104) never contributes to a real push-down, and dropping it would shrink the surface the rule must be trusted over. If you do have a firing case in mind, could you add a test? The resultExpressions rename path and the WindowGroupLimitExec traversal are both currently uncovered.

[ulysses-you]

Agreed, and thanks for the careful trace — your reasoning is exactly right. I verified it three ways (constructed plans, a plan-level probe with the window and aggregate co-partitioned so no shuffle separates them, and the general argument): to push through a SortAggregateExec, line 118 forces the ordering to prefix-cover the grouping keys and line 119 forces every key to be a child column, but the aggregate’s child is already sorted by the full grouping keys, and any extra column would be an aggregate result absent from child.outputSet. So the pushed ordering can never be strictly wider than the sort EnsureRequirements places directly beneath the aggregate — the branch never fires.

Dropped SortAggregateExec from isOrderPreserving and removed the resultExpressions rename path. The canonical window-to-sort-aggregate test still passes: there the aggregate is the top consumer and the push-down traverses the Window, so the aggregate itself is never traversed.

[ulysses-you]

Thanks for the thorough review @peter-toth! Addressed all three in 2b05aa93ee0:

1. Window order-sensitivity — kept on purpose (only refines tie order for a non-unique window ORDER BY, already non-deterministic), documented the distinction from CollectMetricsExec at the isOrderPreserving entries, and softened the PR description.
2. Window-to-sort-aggregate test — the old data had c constant within (a, b); switched to id % 13 + tie-safe RANK() so the widening genuinely reorders window ties while the result stays stable.
3. SortAggregateExec branch — confirmed unreachable and removed it (and the resultExpressions rename path); the canonical case traverses the Window, not the aggregate.

Replied inline with details on each.

[sunchao]

Summary

Thanks for the PR. I reviewed exact head 2b05aa93ee0e709b198ca9ec81dffc2a370cc01a. The central optimization is useful and the prefix/order-availability checks are generally well structured, but I found two issues inline: moving a sort across nondeterministic projects/filters can change row-value associations or survivors, and moving the wider sort below a cardinality reducer can introduce a large default-on performance regression.

Prior state and problem

EnsureRequirements can produce multiple local sorts in one stage when stacked operators require prefix-compatible orderings. RemoveRedundantSorts cannot combine these when the wider ordering is above the narrower one, so the stage retains avoidable sorting work.

Design approach

The new physical rule starts at a local SortExec, walks through selected order-preserving unary operators, rewrites ordering attributes through plain aliases, and widens a strictly narrower lower local sort. It is installed after EnsureRequirements for both regular execution and AQE and is protected by an internal default-on SQL configuration.

Correctness / compatibility analysis

The strict-cover test, sort-key availability check, preservation of direction/null ordering, shuffle boundary, and plain-alias rewrite all look sound. Derived aliases are correctly rejected when their attributes are unavailable below the lower sort. The blocking semantic issue is that ProjectExec and FilterExec are crossed without the determinism guards used by Spark's existing logical sort-elimination rule.

Key design decisions

The rule deliberately handles only strict ordering extensions and leaves equivalent-sort cleanup to RemoveRedundantSorts. It does not cross global sorts or non-order-preserving operators. Refinement of non-unique window tie order is documented as an accepted nondeterministic behavior, distinct from the deterministic-expression problem called out inline.

Implementation sketch

This patch adds the SQL configuration, the PushDownLocalSort physical rule, regular and adaptive preparation hooks, shared AQE/non-AQE SQL tests, and direct physical-plan tests for traversal and alias rewriting.

Behavioral changes worth calling out

Besides removing sorts, the rule changes where comparison work occurs. Crossing a selective filter or window group limit can move the wider sort from a small survivor set to the full input, so a syntactic reduction in sort count is not necessarily a runtime improvement.

Suggested improvements

Please add determinism checks for project/filter traversal and avoid crossing cardinality reducers without a cost proof. Regression coverage should compare exact enabled/disabled results for seeded nondeterministic expressions and include negative plans for selective filters and WindowGroupLimitExec.

Validation: git diff --check passed, and under Java 17 both PushDownLocalSortSuite variants passed: 24 tests, 2 suites, 0 failures.

[sunchao]

[P1] Keep the sort above nondeterministic projects and filters

These two cases are treated as order-preserving regardless of expression determinism. A reachable shape is Sort[a,b] -> Project[a,b,rand(0)] -> Sort[a]: with this rule disabled, the seeded random stream is evaluated after the narrow sort; with it enabled, the [a,b] sort moves below the project, so different rows consume each random value. The analogous Filter(rand(0) < 0.5) shape can retain a different set of rows.

This is the same semantic hazard that EliminateSorts.canEliminateSort guards with p.projectList.forall(_.deterministic) and f.condition.deterministic. Please mirror those checks here (or exclude FilterExec as suggested separately) and add enabled/disabled result tests using a seeded nondeterministic expression.

[ulysses-you]

Great catch — reproduced it (Sort[a,b] -> Project[a,b,rand(0)] -> Sort[a] did push the sort below the project). Fixed by mirroring EliminateSorts.canEliminateSort: ProjectExec is now gated on projectList.forall(_.deterministic) and FilterExec on condition.deterministic. Added plan-level tests that a non-deterministic project and a non-deterministic filter are both left unchanged.

[sunchao]

[P2] Avoid widening the large-input sort below a selective reducer

Crossing FilterExec here (and WindowGroupLimitExec below) can reverse the intended cost improvement. For stacked windows, the disabled plan can sort N rows by narrow [a,b], reduce to R rows, and sort only those survivors by [a,b,c]; this rewrite instead sorts all N rows by [a,b,c] before the reduction. With rn = 1 or a selective rank predicate, R << N, and suffix width plus reduction ratio can make the default-on regression arbitrarily large relative to the removed survivor sort.

Please initially stop at FilterExec and WindowGroupLimitExec, or require a reliable cardinality/cost proof before crossing them. Negative plan tests for a selective predicate and rn = 1 would cover both paths.

[ulysses-you]

Agreed. I made crossing the cardinality reducers opt-in: FilterExec and WindowGroupLimitExec are now only traversed when a new internal config spark.sql.execution.pushDownLocalSort.throughCardinalityReducer is enabled (default false), so by default the rule stops at them and never moves the wider sort below a selective reducer. Added tests that a filter is not crossed by default and is crossed only when the config is on.

[ulysses-you]

Thanks @sunchao! Addressed both in dd42f7318f7:

• [P1] Non-deterministic project/filter — reproduced (the sort was pushed below Project[..., rand(0)]). Now gated on determinism, mirroring EliminateSorts.canEliminateSort: ProjectExec on projectList.forall(_.deterministic), FilterExec on condition.deterministic.
• [P2] Cardinality reducers — made crossing FilterExec/WindowGroupLimitExec opt-in via a new internal config spark.sql.execution.pushDownLocalSort.throughCardinalityReducer (default false), so by default the rule never moves the wider sort below a selective reducer.

Added plan-level tests for both (non-deterministic project/filter left unchanged; filter crossed only when the reducer config is on). Replied inline with details.

[sunchao]

Thanks for the fixes. I re-reviewed exact head 4b5520cae9e98b445b412e0619a1cc4f6033211f with five independent passes. The earlier [P1] determinism issue is fixed by the ProjectExec/FilterExec guards, and the earlier [P2] cost issue is fixed by keeping cardinality-reducer traversal opt-in and default-off. The focused AQE/non-AQE suites pass 30/30 locally, so I’m approving.

Two non-blocking follow-ups:

• [P3] Sync the PR description with the final code. It still lists SortAggregateExec as traversed and presents FilterExec/WindowGroupLimitExec without the new default-off reducer gate.
• [P3] Strengthen the AQE assertion. checkNumSorts inspects executedPlan before collect(). Since this rule also runs in AQE query-stage preparation and replanning, consider executing first and checking the final adaptive plan as well.

The existing WindowGroupLimitExec coverage suggestion is already captured in the earlier review thread, so I’m not duplicating it here.

[ulysses-you]

Thanks for the approval @sunchao! Addressed both non-blocking [P3] follow-ups:

• Synced the PR description with the final code — removed SortAggregateExec from the traversed operators, and documented that FilterExec/WindowGroupLimitExec are crossed only when spark.sql.execution.pushDownLocalSort.throughCardinalityReducer (default false) is enabled, plus the determinism guard on project/filter.
• Strengthened the AQE assertion (b6503d817a0) — checkNumSorts now executes the query before inspecting the plan, so under AQE it checks the final adaptive plan (after query-stage materialization/replanning) rather than the initial plan.

The WindowGroupLimitExec coverage suggestion from the earlier thread: it is now gated behind the reducer config and covered indirectly by the reducer-gating test; happy to add a dedicated WindowGroupLimitExec traversal test if you think it is worth it.

[peter-toth]

Thanks for addressing the review, @ulysses-you — re-checked through b6503d817a0; my findings 1-3 (window order-sensitivity documented, the neutralized window test, the dead SortAggregateExec branch) are all resolved, nothing new from me.