jayzhan211 commented on code in PR #26459:
URL: https://github.com/apache/datafusion/pull/26459#discussion_r4324942456
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs:
##########
@@ -232,6 +235,71 @@ pub fn replace_with_partial_sort(
Ok(plan)
}
+/// Adds a [`SortExec`] below every node whose ordering requirement is not
+/// satisfied by its child's equivalence properties.
+///
+/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a
+/// descendant that satisfies it. That mapping is computed separately from the
+/// equivalence properties of the operators in between, and the two can
+/// disagree. The plan may then be correct, but its properties cannot show it,
+/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan
+/// produced by `EnsureRequirements` valid by the same check.
+///
+/// A requirement that a new [`SortExec`] still cannot satisfy is left as is,
+/// for `SanityCheckPlan` to report.
+///
+/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts
+pub fn add_sort_for_unproven_ordering(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ let requirements = plan.required_input_ordering();
+ if requirements.iter().all(Option::is_none) {
+ return Ok(Transformed::no(plan));
+ }
+ let mut children: Vec<_> = plan.children().into_iter().cloned().collect();
+ let mut changed = false;
+ for (child, requirement) in children.iter_mut().zip(requirements) {
+ let Some(requirement) = requirement else {
+ continue;
+ };
+ let requirement = requirement.into_single();
+ let eq_properties = child.equivalence_properties();
+ if eq_properties.ordering_satisfy_requirement(requirement.clone())? {
+ continue;
+ }
+ let sort_reqs: Vec<_> = requirement.clone().into();
+ let sort_exprs = sort_reqs
+ .into_iter()
+ .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none())
+ .map(Into::into);
+ let Some(ordering) = LexOrdering::new(sort_exprs) else {
+ continue;
+ };
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(
+ SortExec::new(ordering,
Arc::clone(child)).with_preserve_partitioning(
+ child.output_partitioning().partition_count() > 1,
+ ),
+ );
+ if !sort
+ .equivalence_properties()
+ .ordering_satisfy_requirement(requirement)?
+ {
+ continue;
+ }
+ *child = sort;
+ changed = true;
+ }
+ if changed {
+ plan.replace_children(
Review Comment:
This no longer applies: following your comment on the root cause, the
EnsureRequirements step is removed and the PR fixes the join instead.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs:
##########
@@ -232,6 +235,71 @@ pub fn replace_with_partial_sort(
Ok(plan)
}
+/// Adds a [`SortExec`] below every node whose ordering requirement is not
+/// satisfied by its child's equivalence properties.
+///
+/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a
+/// descendant that satisfies it. That mapping is computed separately from the
+/// equivalence properties of the operators in between, and the two can
+/// disagree. The plan may then be correct, but its properties cannot show it,
+/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan
+/// produced by `EnsureRequirements` valid by the same check.
+///
+/// A requirement that a new [`SortExec`] still cannot satisfy is left as is,
+/// for `SanityCheckPlan` to report.
+///
+/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts
+pub fn add_sort_for_unproven_ordering(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ let requirements = plan.required_input_ordering();
+ if requirements.iter().all(Option::is_none) {
+ return Ok(Transformed::no(plan));
+ }
+ let mut children: Vec<_> = plan.children().into_iter().cloned().collect();
+ let mut changed = false;
+ for (child, requirement) in children.iter_mut().zip(requirements) {
+ let Some(requirement) = requirement else {
+ continue;
+ };
+ let requirement = requirement.into_single();
+ let eq_properties = child.equivalence_properties();
+ if eq_properties.ordering_satisfy_requirement(requirement.clone())? {
+ continue;
+ }
+ let sort_reqs: Vec<_> = requirement.clone().into();
+ let sort_exprs = sort_reqs
+ .into_iter()
+ .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none())
+ .map(Into::into);
+ let Some(ordering) = LexOrdering::new(sort_exprs) else {
+ continue;
+ };
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(
Review Comment:
This no longer applies: following your comment on the root cause, the
EnsureRequirements step is removed and the PR fixes the join instead.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs:
##########
@@ -232,6 +235,71 @@ pub fn replace_with_partial_sort(
Ok(plan)
}
+/// Adds a [`SortExec`] below every node whose ordering requirement is not
+/// satisfied by its child's equivalence properties.
+///
+/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a
+/// descendant that satisfies it. That mapping is computed separately from the
+/// equivalence properties of the operators in between, and the two can
+/// disagree. The plan may then be correct, but its properties cannot show it,
+/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan
+/// produced by `EnsureRequirements` valid by the same check.
+///
+/// A requirement that a new [`SortExec`] still cannot satisfy is left as is,
+/// for `SanityCheckPlan` to report.
+///
+/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts
+pub fn add_sort_for_unproven_ordering(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ let requirements = plan.required_input_ordering();
+ if requirements.iter().all(Option::is_none) {
+ return Ok(Transformed::no(plan));
+ }
+ let mut children: Vec<_> = plan.children().into_iter().cloned().collect();
+ let mut changed = false;
+ for (child, requirement) in children.iter_mut().zip(requirements) {
+ let Some(requirement) = requirement else {
+ continue;
+ };
+ let requirement = requirement.into_single();
+ let eq_properties = child.equivalence_properties();
+ if eq_properties.ordering_satisfy_requirement(requirement.clone())? {
+ continue;
+ }
+ let sort_reqs: Vec<_> = requirement.clone().into();
+ let sort_exprs = sort_reqs
+ .into_iter()
+ .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none())
+ .map(Into::into);
+ let Some(ordering) = LexOrdering::new(sort_exprs) else {
+ continue;
+ };
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(
+ SortExec::new(ordering,
Arc::clone(child)).with_preserve_partitioning(
Review Comment:
This no longer applies, since the step is removed. For the record,
`LimitPushdown` ran after it and turned the added sort into a TopK under a
`LIMIT`.
--
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]
