xudong963 commented on code in PR #23301:
URL: https://github.com/apache/datafusion/pull/23301#discussion_r3518021323
##########
datafusion/physical-optimizer/src/pushdown_sort.rs:
##########
@@ -172,13 +182,38 @@ impl PhysicalOptimizerRule for PushdownSort {
// Data source is optimized for the ordering but not
perfectly sorted
// Keep the Sort operator but use the optimized input
// Benefits: TopK queries can terminate early, better
cache locality
- Ok(Transformed::yes(Arc::new(
+ // A standalone multi-partition TopK still needs a global
+ // merge; otherwise a later coalesce can concatenate
+ // locally sorted partitions.
+ let preserve_partitioning =
sort_exec.preserve_partitioning();
+ let needs_global_topk =
+ sort_exec.preserve_partitioning() &&
sort_exec.fetch().is_some();
+ let input_partitions = sort_input
+ .output_partitioning()
+ .partition_count()
+ .max(inner.output_partitioning().partition_count());
Review Comment:
Done, changed this to use inner.output_partitioning().partition_count()
directly.
##########
datafusion/physical-optimizer/src/pushdown_sort.rs:
##########
@@ -172,13 +182,38 @@ impl PhysicalOptimizerRule for PushdownSort {
// Data source is optimized for the ordering but not
perfectly sorted
// Keep the Sort operator but use the optimized input
// Benefits: TopK queries can terminate early, better
cache locality
- Ok(Transformed::yes(Arc::new(
+ // A standalone multi-partition TopK still needs a global
+ // merge; otherwise a later coalesce can concatenate
+ // locally sorted partitions.
+ let preserve_partitioning =
sort_exec.preserve_partitioning();
+ let needs_global_topk =
+ sort_exec.preserve_partitioning() &&
sort_exec.fetch().is_some();
Review Comment:
Done, reused the preserve_partitioning local variable for needs_global_topk.
##########
datafusion/core/tests/physical_optimizer/pushdown_sort.rs:
##########
@@ -119,6 +127,67 @@ fn test_sort_with_limit_phase1() {
);
}
+#[test]
Review Comment:
Done, added a single-partition negative regression test that keeps the
standalone TopK without adding SortPreservingMergeExec.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
