[yashrb24]

Which issue does this PR close?

Related to #25033.

Rationale for this change

Sorting evaluates all ORDER BY expressions up front, even when earlier keys already determine the order. This can waste work when computing later expressions are expensive.

This PR introduces staged evaluation so later expressions are evaluated only for rows still tied on earlier keys.

What changes are included in this PR?

Select stages automatically, grouping cheap keys and starting new stages before expensive expressions. This is combined with a heuristic where we determine further staging is not useful and fallback to the upfront sorting of the remaining keys.

This can also be extended TopK and spill sorting but have kept them outside the PR's scope for now.

What is the testing strategy for this PR?

Some new tests have been added on this specific behaviour

Are there any user-facing changes?

Adds datafusion.execution.enable_staged_sort, defaulting to false.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 91.52542% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.84%. Comparing base (22651d2) to head (094d0a2).
⚠️ Report is 384 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/sorts/sort.rs	91.44%	7 Missing and 28 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25358      +/-   ##
==========================================
+ Coverage   81.92%   82.84%   +0.91%     
==========================================
  Files        1135     1149      +14     
  Lines      427772   452993   +25221     
  Branches   427772   452993   +25221     
==========================================
+ Hits       350456   375281   +24825     
+ Misses      56373    54868    -1505     
- Partials    20943    22844    +1901     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[kosiew]

@yashrb24
Please amend the PR description to follow the .github/pull_request_template.md template.

[kosiew]

@yashrb24,

Thanks for working on this. I found two issues that I think need to be addressed before this is ready to merge.

[kosiew]

Staged sorting can skip later fallible or volatile ORDER BY expressions when earlier keys already resolve the order, so queries can behave differently depending on tie distribution. Please either document this experimental/error/volatile-evaluation contract and add tests that pin it, including that disabled staging stays eager, or preserve eager evaluation.

[yashrb24]

thank you. Though I wanted some directional feedback before actually having the full fledged implementation out and accounting for all cases. If the idea for this PR looks good I'll have these fixed too

[kosiew]

The new staged sorting path does not have focused behavioral tests. Please add coverage for staged versus eager equivalence, successive tie groups, null ordering, chunked output, the zero-tie/error case, and both single-batch and multiple in-memory-run paths.

[yashrb24]

thanks for the catch, have added them

[kosiew]

@yashrb24,

Thanks for working on staged sorting. I think this is a promising direction, and the benchmarks show substantial benefits when earlier keys resolve most rows.

I found a correctness regression with column-free expressions in later stages. Please fix this and add regression coverage before merging.

The evaluation contract also needs to document that staged sorting can skip fallible or volatile expressions. Please clarify this behavior and its implications.

I also recommend adding reproducible benchmarks covering distinct keys, small and large tie groups, all-tied inputs, and multi-batch sorting. In a local micro-benchmark, 500k groups of two rows caused a 1.51x slowdown with a cheap CAST expression. Please investigate the per-group sorting and allocation overhead.

The optimization currently excludes TopK and does not avoid expression reevaluation during merging. It would help to clarify whether these are planned follow-ups.

Overall, I support continuing with the opt-in approach, but the correctness regression and evaluation-contract gap should be addressed before merging.

[kosiew]

When a later stage uses a column-free expression, stage.columns is empty and take_record_batch fails because the empty-column batch loses its row count. Could we preserve the selected row count using RecordBatchOptions::with_row_count(Some(selection.len())) or retain the full-batch gather for this case, and add tests with a zero-argument UDF and tied keys for both single and concatenated batches?

[kosiew]

@yashrb24

I think manual stage grouping does not seem like a practical production setting. The optimal group size depends on expression cost and tie distribution, which vary across queries and batches. With the default sort_key_group_size = 5, enabling staged sorting has no effect on typical two-key queries.

I suggest exploring three ways to automate this (all unvalidated and requiring benchmarks):

1. Split by key cost: Group cheap keys together and start new stages before expensive expressions such as UDFs.
2. Adapt at runtime: Use tie counts after each stage to decide whether to continue staging or evaluate the remaining keys together, avoiding overhead from many tiny tie groups.
3. Simplify configuration: Keep at most an on/off switch and remove sort_key_group_size. This would also eliminate one of the two new public configuration fields contributing to the semver break.

I think automatic staging is a better long-term direction than requiring users to tune group sizes manually.

[yashrb24]

@kosiew thanks for the feedback and sorry for the delayed response!
So a bunch of observations from some of my attempts to do automatic staging:

1. A lot of the numbers depend on what one categorizes as a "cheap key" and how we define a stage. We do have some precedence for defining this set of cheap keys, but it exists in the logical planning step and not sure if we want to reproduce all of that infra here too. If this is the direction you are aligned on, I'll work on this.
2. One of the drawbacks with this approach of "split by key cost" that I can see is that it will potentially not help in cases where we have multiple expensive order-by expressions stacked together (eg: order by expr(col1), expr(col2) ... expr(coln) where expr is some expensive operation), especially when we have ties. Limiting this overhead was why originally the reason of going with the fixed grouped size. Again a lot of this will depend on the distribution of data as you have highlighted. For this, the "Adapt at runtime" is a useful fallback, I'll have this implemented.

[kosiew]

Converting to draft while awaiting changes

[yashrb24]

@kosiew I have updated the PR based on your feedback. Please take another look!

Even with automatic staging, the improvements are still dependent on expression cost and tie distribution. And can potentially cause degradation from existing logic because of the overhead from staging in some cases. But the motivation is that earlier ORDER BY keys can determine the order without evaluating every later expression.

I've updated the relevant context section of #25033 for some sources on this