xudong963 commented on code in PR #24809:
URL: https://github.com/apache/datafusion/pull/24809#discussion_r3953995455
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1722,9 +1756,26 @@ pub fn ensure_distribution(
replace_children_if_necessary(plan, children_plans)?
};
- Ok(Transformed::yes(DistributionContext::new(
- plan, data, children,
- )))
+ let mut optimized_context = DistributionContext::new(plan, data, children);
+
+ // A removed fetch must survive even when this node does not need a new
+ // distribution operator. Otherwise a second optimizer pass can silently
+ // remove the query's LIMIT.
+ if let Some(fetch) = removed_fetch {
Review Comment:
Thanks for the concrete example. Fixed in 8098706: distribution cleanup now
stops at fetched operators, keeping each limit at its original position instead
of collapsing fetch values and restoring only the outermost operator. Replacing
a fetched SPM also retains its TopK at that position with the original merge
ordering, even when the parent requires a different ordering.
The nested execution regression uses even/odd partitions and outer fetch
values of 0, 3, and 10 around an SPM fetching 5. It checks exact results before
optimization and after two passes, with sort parallelization both enabled and
disabled.
I also protected fetched Coalesce nodes from removal during sort
parallelization and replacement with ordered merges. Added regressions for
reversed ancestor ordering and a filter above a fetched coalesce. All three new
execution regressions fail on c90e58b and pass with this fix.
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -4558,6 +4559,166 @@ fn test_replace_order_preserving_variants_with_fetch()
-> Result<()> {
Ok(())
}
+#[test]
+fn preserve_fetch_when_reoptimizing_ordered_merge() -> Result<()> {
+ let schema = schema();
+ let sort_key: LexOrdering =
+ [PhysicalSortExpr::new_default(col("c", &schema)?)].into();
+ let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]);
+ let plan: Arc<dyn ExecutionPlan> =
+ Arc::new(SortPreservingMergeExec::new(sort_key,
input).with_fetch(Some(5)));
+
+ let optimized =
+ EnsureRequirements::new().optimize(plan,
&test_suite_default_config_options())?;
+ let plan = displayable(optimized.as_ref()).indent(true).to_string();
+
+ assert!(
+ plan.contains("SortPreservingMergeExec: [c@2 ASC], fetch=5"),
+ "expected the optimizer to preserve fetch:\n{plan}"
+ );
+
+ Ok(())
+}
+
+#[test]
+fn preserve_fetch_when_reoptimizing_coalesce_partitions() -> Result<()> {
+ let input = parquet_exec_multiple();
+ let plan: Arc<dyn ExecutionPlan> =
+ Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5)));
+
+ let optimized =
+ EnsureRequirements::new().optimize(plan,
&test_suite_default_config_options())?;
+
+ assert_eq!(optimized.fetch(), Some(5));
+ optimized
+ .downcast_ref::<CoalescePartitionsExec>()
+ .expect("expected CoalescePartitionsExec");
+
+ Ok(())
+}
+
+#[test]
+fn move_fetch_to_replacement_sort() -> Result<()> {
Review Comment:
Done in 8098706. `move_fetch_to_replacement_sort` now executes two sorted
partitions and compares the actual TopK values before and after replacement,
while still checking the replacement sort's `fetch=5`.
The cases cover ascending order with NULLs first, descending order with
NULLs last, and duplicate values. This checks the resulting values and ordering
as well as the plan shape.
--
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]
