Repository navigation
Conversation
github-actions
Bot
added
documentation
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
github-actions
Bot
removed
documentation
github-actions
Bot
added
documentation
Codecov Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
|
@yashrb24 |
| } | ||
| arrow::compute::take_record_batch(batch, &UInt32Array::from(selection))? | ||
| }; | ||
| let sort_columns = keys |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
| } | ||
| } | ||
|
|
||
| fn staged_sort_batch_chunked( |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
thanks for the catch, have added them
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
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.
| selection.extend_from_slice(&indices[range.clone()]); | ||
| } | ||
| arrow::compute::take_record_batch( | ||
| &batch.project(&stage.columns)?, |
There was a problem hiding this comment.
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
left a comment
kosiew
left a comment
There was a problem hiding this comment.
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):
- Split by key cost: Group cheap keys together and start new stages before expensive expressions such as UDFs.
- 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.
- 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.
3 participants
Which issue does this PR close?
Related to #25033. This PR is mostly an ideation of the way we can bring in this change since it's kind of an invasive change to the sorting mechanics and thus wanted feedback.
Rationale for this change
Today, sorting evaluates all
ORDER BYexpressions before it compares rows. Some of this work is unnecessary when earlier keys already determine the order.This change evaluates later expressions only for rows tied on earlier keys. This could help when later expressions are expensive and earlier keys resolve most rows. The drawback is that the extra sorting steps can outweigh these savings when expressions are cheap or many rows remain tied.
What changes are included in this POC PR?
Are there any user-facing changes?
had added these flags for my local testing, can remove if needed
datafusion.execution.enable_staged_sort: defaultfalse.datafusion.execution.sort_key_group_size: default value is5Currently not extending to TopK or spill codepath