[AURON #2463] Fix native Top-K with only sort keys - #2464
Conversation
There was a problem hiding this comment.
Pull request overview
Fixes a native Top-K failure in SortExec when all output columns are also sort keys, by ensuring the zero-column RecordBatch created for the output uses the retained (post-limit) row count so downstream merge cursors don’t read past available sort-key rows.
Changes:
- Adjust zero-column batch row count to
sorted_indices.len()after Top-K limiting in the native external sorter. - Add a Spark regression test covering
ORDER BY ... LIMITwhen projecting only the sort column. - Add a DataFusion-native regression test for Top-K with a single (sort) column.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| spark-extension-shims-spark/src/test/scala/org/apache/auron/exec/AuronExecSuite.scala | Adds a Spark-side regression test ensuring native TakeOrdered works when projecting only the sort column. |
| native-engine/datafusion-ext-plans/src/sort_exec.rs | Fixes zero-column batch row count for Top-K and adds a native regression test for the scenario. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| let expected = r#"+----+ | ||
| | id | | ||
| +----+ | ||
| | 0 | | ||
| | 1 | | ||
| | 2 | | ||
| | 3 | | ||
| | 4 | | ||
| | 5 | | ||
| +----+"#; | ||
| assert_batches_eq!(expected.lines().collect::<Vec<_>>(), &batches); |
There was a problem hiding this comment.
rustfmt compacts the suggested Vec<&str> form into a few long lines, making the expected table less readable. The raw string has no leading spaces, so I prefer to keep the current format.
| let sorted_batch = if !self.prune_sort_keys_from_batch.is_all_pruned() { | ||
| take_batch(batch, sorted_indices)? | ||
| } else { | ||
| create_zero_column_batch(batch.num_rows()) | ||
| create_zero_column_batch(sorted_indices.len()) | ||
| }; |
There was a problem hiding this comment.
The observation about over-allocation is correct, but this allocation behavior predates this pr and remains unchanged by it. Changing the reservation strategy is independent of this correctness fix and requires separate validation, so it is better suited to a followup pr.
slfan1989
left a comment
slfan1989
left a comment
There was a problem hiding this comment.
Thanks for the fix. Using the retained sorted-index count keeps the zero-column batch aligned with the separately stored sort keys, and the Rust and Spark regression tests cover the failing Top-K path well. LGTM.
|
@lyne7-sc Thanks for the contribution! Merged into master. |
|
Thanks for the reviews! @slfan1989 @SteNicholas |
Which issue does this PR close?
Closes #2463
Rationale for this change
When all output columns are sort keys, the native sorter stores the sort keys separately and creates a zero-column data batch.
For Top-K queries, the sorted indices are limited before this batch is created, but the batch currently retains the original input row count. This mismatch can cause the merge cursor to read past the retained sort-key rows.
What changes are included in this PR?
Are there any user-facing changes?
Bug fix only.
How was this patch tested?
Added a regression test to AuronExecSuite.
Was this patch authored or co-authored using generative AI tooling?
If yes, include:
Generated-by: GPT-5ASF guidance: https://www.apache.org/legal/generative-tooling.html